# TensorFlow and tf.keras
#import tensorflow as tf#dont need it here;)

# Helper libraries
import numpy as np
import matplotlib.pyplot as plt
import math

from xevo import eobj, quickmut,semiquickmut
from ising2d import ising2devo

import os
from onceasecond import onceasecond
from time import time

from data import loadfile
from sklearn.metrics import roc_curve,auc

def cauc(n,a):
    p=np.concatenate((n,a),axis=0)
    l=np.concatenate((np.zeros(len(n)),np.ones(len(a))),axis=0)
    fpr, tpr, threshold = roc_curve(l, p)

    auc_score=auc(fpr,tpr)
    return auc_score



class matrixject(eobj):

    def norm(s):
        return np.mean(s.q**2)*len(s.q)

    def __init__(s,mat=None,dim=2):

        if mat is None:
            s.q=np.random.normal(0,1,(dim,dim))
            s.q/=s.norm()+0.00001
        else:
            s.q=mat


        s.initial()

    def __add__(a,b):
        return a.__class__((a.q+b.q)/2)
    def randomize(s):
        s=s.copy()
        s.q=np.random.normal(0,1,s.q.shape)
        s.q/=s.norm()+0.00001
        return s

    def mutate(s):
        s=s.copy()
        i1,i2=np.random.randint(0,len(s.q),2)
        val=np.random.normal(0,1)
        exp=np.random.random()*5-3
        s.q[i1,i2]+=val*exp
        return s
    def predict_data(s,data):
        mat=(s.q-1)
        pred=np.dot(data,mat)
        return pred
    def eval_data(s,data):
        return np.mean(np.abs(s.predict_data(data)),axis=-1)
    def calcstrength(s):
        #mat*x=x
        #(mat-1)*x=0
        mat=(s.q-1)
        loss=np.mean(s.eval_data(data))
        #loss+=1/(1+np.mean(np.abs(mat)))#punish trivial solution
        return loss

    def shallmaximize(s):
        return False

    def _copy(s):
        return s.__class__(s.q)

    def __repr__(s):
        return str(s.q)


class ising2dlog(ising2devo):
    def __init__(s,*args,**kwargs):
        ising2devo.__init__(s,*args,**kwargs)
        s.starttime=time()
    @onceasecond
    def stepsave(s):
        return #disabled here
        np.savez_compressed(logfile.replace("log.csv",f"step{s.i}"),mats=[zw.q for zw in s.q],stres=[zw.strength() for zw in s.q],i=s.i)
    def logeneration(s,show=True):
        mx,mn,std=ising2devo.logeneration(s=s,show=show)
        global logfile
        with open(logfile,"a") as f:
            f.write(f"{mx},{mn},{std},")
            f.write(",".join([str(qq.strength()) for qq in s.q]))
            f.write("\n")
        s.stepsave()
        t=time()
        if t-s.starttime>600:#10min
            mx=-1.0
        return mx,mn,std



def train(x,y,pth,typ):
    x=x.astype("float32")
    y=y.astype("float32")

    x/=np.mean(np.abs(x))

    global logfile
    logfile=pth+"/log.csv"

    #x-=np.mean(x,axis=0)#no correcting here
    #x/=(np.std(x,axis=0)+0.00001)

    train=x[np.where(y[:,0]==0)]
    test_a=x[np.where(y[:,0]==1)]
    test_n=train[:len(test_a)]
    train=train[len(test_a):]

    #global dd
    #dd=train

    #exit()

    global data
    data=train


    obj=matrixject(dim=int(data.shape[1]))

    boardsize=10

    opt=ising2dlog(m1=boardsize,m2=boardsize,#dimension of the board, resulting in population of 49
            temp=10.0,#semi arbitrary hyper param
            mergews=0.3,#high merge ws since (1+1)/2=1
            close="c")#classical ising neighbourhood



    solv=semiquickmut(obj,goal=0.0,maxsteps=500000,population=boardsize*boardsize,opt=opt)

    stres=[zw.strength() for zw in solv]
    matrices=[zw.q for zw in solv]
    sol=solv[np.argmin(stres)]


    alle=[zw.eval_data(train) for zw in solv]
    allen=[zw.eval_data(test_n) for zw in solv]
    allea=[zw.eval_data(test_a) for zw in solv]

    lss = sol.calcstrength()
    #p=sol.predict_data(train)
    #pn=sol.predict_data(test_n)
    #pa=sol.predict_data(test_a)

    e=sol.eval_data(train)
    en=sol.eval_data(test_n)
    ea=sol.eval_data(test_a)

    mn=np.mean(e)
    dn=(en-mn)**2
    da=(ea-mn)**2

    auc=cauc(dn,da)
    print(f"auc={auc}")



if __name__ == '__main__':
    x,y=loadfile()
    train(x,y,".","")