# -*- coding: utf-8 -*-

import torch
from codebert import codebert_mlm, codebert_cls, codebert_clone
from codegpt import codegpt
import torch.nn.functional as F

import argparse
import os

if __name__ == "__main__":
    
    parser = argparse.ArgumentParser()
    parser.add_argument('--gpu', type=str, default='-1')
    # microsoft/codebert-base-mlm
    parser.add_argument('--mlm_path', type=str,
                        default="/var/data/lushuai/bertvsbert/save/with_rtd/poj/checkpoint-23000-1.0541",
                        help="Path to the masked language model")
    parser.add_argument('--cls_path', type=str,
                        default="/var/data/lushuai/bertvsbert/save/poj-classifier/checkpoint-51000-0.986",
                        help="Path to the OJ classifier")
    parser.add_argument('--clone_path', type=str,
                        default="/var/data/lushuai/bertvsbert/save/poj-clone/old/rtd/checkpoint-best-f1",
                        help="Path to the clone dectection on bigclonebench")
    parser.add_argument('--gpt_path', type=str,
                        default="microsoft/CodeGPT-small-java-adaptedGPT2",
                        help="Path to the codegpt")
    
    opt = parser.parse_args()

    if int(opt.gpu) < 0:
        device = torch.device("cpu")
    else:
        os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpu
        device = torch.device("cuda")
    mlm_model = codebert_mlm(opt.mlm_path, device)
    # cls_model = codebert_cls(opt.cls_path, device)
    clone_model = codebert_clone(opt.clone_path, device, block_size=400)
    # codegpt_model = codegpt(opt.gpt_path, device, block_size=512)
    
    inputs = [
        "int main ( ) { int n , i ; n = 1 ; return 0 }",
        "int main ( ) { int <mask>, i ; <mask> = 1 ; return 0 }", 
        "void main ( ) { double x ; }",
        "int main ( ) { int aVeryLongIntegerVar = 0 ; return aVeryLongIntegerVar ; }",
        "int main ( ) { int n , i ; n = 1 ; return 0 }",
        "int main ( ) { int <mask>, i ; <mask> = 1 ; return 0 }", 
        "void main ( ) { double x ; }",
        "int main ( ) { int aVeryLongIntegerVar = 0 ; return aVeryLongIntegerVar ; }",
        "int main ( ) { int n , i ; n = 1 ; return 0 }",
        "int main ( ) { int <mask>, i ; <mask> = 1 ; return 0 }", 
        "void main ( ) { double x ; }",
        "int main ( ) { int aVeryLongIntegerVar = 0 ; return aVeryLongIntegerVar ; }",
        "int main ( ) { int n , i ; n = 1 ; return 0 }",
        "int main ( ) { int <mask>, i ; <mask> = 1 ; return 0 }", 
        "void main ( ) { double x ; }",
        "int main ( ) { int aVeryLongIntegerVar = 0 ; return aVeryLongIntegerVar ; }",
        "int main ( ) { int n , i ; n = 1 ; return 0 }",
        "int main ( ) { int <mask>, i ; <mask> = 1 ; return 0 }", 
        "void main ( ) { double x ; }",
        "int main ( ) { int aVeryLongIntegerVar = 0 ; return aVeryLongIntegerVar ; }",
        "int main ( ) { int n , i ; n = 1 ; return 0 }",
        "int main ( ) { int <mask>, i ; <mask> = 1 ; return 0 }", 
        "void main ( ) { double x ; }",
        "int main ( ) { int aVeryLongIntegerVar = 0 ; return aVeryLongIntegerVar ; }",
        "int main ( ) { int n , i ; n = 1 ; return 0 }",
        "int main ( ) { int <mask>, i ; <mask> = 1 ; return 0 }", 
        "void main ( ) { double x ; }",
        "int main ( ) { int aVeryLongIntegerVar = 0 ; return aVeryLongIntegerVar ; }"
    ]

    def test_codebert():
        clone_inputs = [
            'int main(){ cout << "Hello World"; return 0;}',
            'void main() { ;}'
        ]

        print(clone_inputs[0])
        print(clone_inputs[1])
        # exit()
        true_inputs = [" ".join(x.split()) for x in clone_inputs]
        # print(true_inputs)
        # exit()
        
        # tokens = cls_model.tokenize(inputs)
        # print (tokens)
        # tokens = mlm_model.tokenize(inputs)
        # print (tokens)
        # pred = cls_model.run(inputs)
        # print (pred.size())
        # print (tokens)
        # pred = mlm_model.run(inputs, len(inputs))
        # print (pred.size())
        # print (tokens)

        # tokens = clone_model.tokenize(clone_inputs)
        # print(tokens)
        # tokens = mlm_model.tokenize(clone_inputs)
        # print(tokens)
        pred = clone_model.run(true_inputs)
        print(pred)
        prob=F.softmax(pred)
        print(prob)
        # pred = mlm_model.run(clone_inputs, len(clone_inputs))
        # print(pred.size())

    def test_codegpt():
        inputs = [
            '    public String[][] getProjecttreeData() {\n        String[][] treeData = null;\n        String filename = dms_home + FS + "temp" + FS + username + "adminprojects.xml";\n        String urlString = dms_url + "/servlet/com.ufnasoft.dms.server.ServerGetAdminProjects";\n        try {\n            String urldata = urlString + "?username=" + URLEncoder.encode(username, "UTF-8") + "&key=" + URLEncoder.encode(key, "UTF-8") + "&filename=" + URLEncoder.encode(username, "UTF-8") + "adminprojects.xml";\n            DocumentBuilderFactory factory = DocumentBuilderFactory.newInstance();\n            factory.setValidating(false);\n            DocumentBuilder parser = factory.newDocumentBuilder();\n            URL u = new URL(urldata);\n            DataInputStream is = new DataInputStream(u.openStream());\n            FileOutputStream os = new FileOutputStream(filename);\n            int iBufSize = is.available();\n            byte inBuf[] = new byte[20000 * 1024];\n            int iNumRead;\n            while ((iNumRead = is.read(inBuf, 0, iBufSize)) > 0) os.write(inBuf, 0, iNumRead);\n            os.close();\n            is.close();\n            File f = new File(filename);\n            InputStream inputstream = new FileInputStream(f);\n            Document document = parser.parse(inputstream);\n            NodeList nodelist = document.getElementsByTagName("proj");\n            int num = nodelist.getLength();\n            treeData = new String[num][3];\n            for (int i = 0; i < num; i++) {\n                treeData[i][0] = new String(DOMUtil.getSimpleElementText((Element) nodelist.item(i), "pid"));\n                treeData[i][1] = new String(DOMUtil.getSimpleElementText((Element) nodelist.item(i), "ppid"));\n                treeData[i][2] = new String(DOMUtil.getSimpleElementText((Element) nodelist.item(i), "p"));\n            }\n        } catch (MalformedURLException ex) {\n            System.out.println(ex);\n        } catch (ParserConfigurationException ex) {\n            System.out.println(ex);\n        } catch (NullPointerException e) {\n        } catch (Exception ex) {\n            System.out.println(ex);\n        }\n        return treeData;\n    }\n',
            '    private static void copy(String sourceName, String destName) throws IOException {\n        File source = new File(sourceName);\n        File dest = new File(destName);\n        FileChannel in = null, out = null;\n        try {\n            in = new FileInputStream(source).getChannel();\n            out = new FileOutputStream(dest).getChannel();\n            long size = in.size();\n            MappedByteBuffer buf = in.map(FileChannel.MapMode.READ_ONLY, 0, size);\n            out.write(buf);\n        } finally {\n            if (in != null) in.close();\n            if (out != null) out.close();\n        }\n    }\n',
        ]
        inputs = [" ".join(x.split()) for x in inputs]
        for x in inputs:
            ppl = codegpt_model.run(x)
            print(ppl)

    # test_codegpt()
    test_codebert()