models =['aae','char_rnn','vae','organ']
#models =['char_rnn','vae'] # organ fails to run

repeats = list(range(3))
trainsizes =[200,400,600,800,1200] # provide these as args so that it can be read later


# this newtraindata was generated by the moses datamaker and then linked to the moses 
# so moses wrote in there.. 
path = 'mosestrain'
import rdkitutils as rdk
import numpy as np
import main2 as main
import sys
models=[main.args.model]

def lc(model='aae',rep=0):
    p= rdk.moses_to_nx(f"{path}/ptest_{rep}.csv")
    n= rdk.moses_to_nx(f"{path}/ntest_{rep}.csv")
    scorer = main.make_scorer(p,n) 
    combined =[]
    base =[]
    newonly=[]
    for size in trainsizes: 
        p = rdk.moses_to_nx(f"{path}/p{size}_{rep}/train.csv")
        n = rdk.moses_to_nx(f"{path}/n{size}_{rep}/train.csv")
        op = list(rdk.smi_to_nx(f"{path}/p{size}_{rep}/{model}/gen"))
        on = list(rdk.smi_to_nx(f"{path}/n{size}_{rep}/{model}/gen"))
        combined.append(scorer(p+op,n+on))
        base.append(scorer(p,n))
        newonly.append(scorer(op,on))
    return combined,base,newonly 

for model in models:
    print (model)
    a,b,c = list(zip(*[lc(model,x) for x in repeats]))
    main.format_abc(a,b,c,f"{model}.pickle") 

