#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu May 28 07:55:23 2020

@author: p277634
"""

import numpy as np
import matplotlib.pyplot as pl
from itertools import combinations
import scipy.io as sio
import copy


class BaseData():
    def __init__(self,dataIni, stages):
        self.nroi   = 68
        self.d      = None
        self.nStg   = len(stages)
        # self.stgIx  = stages-1
        self.sizePop    = np.load(dataIni, allow_pickle=True)['sizeP'].item()
        self.pMig       = np.load(dataIni, allow_pickle=True)['pMig'].item()
        self.nIsla      = np.load(dataIni, allow_pickle=True)['nIsla'].item()
    

class LoadData():
    def __init__(self):
        self.pathOptData    = './data/'
        self.velocity       = 5
        self.nameModel      = 'S1_StgSequence_allPu1000_wE_P030_optIni_v' + str(self.velocity)
        self.nameEnding     = '_fixGL_Stg'
        self.stages         = [0] # list with stages indexes
        self.var2logIx      = {'pulse':[0],'order':[1,2],'fit':[3,5,6],'fc':[4],'env':[7,8,9]}
        
    def getBestModel(self):
        dataIni        = self.pathOptData + self.nameModel + str(self.stages[0]) + self.nameEnding + '.npz'
        m = BaseData(dataIni, self.stages)
        dataread = self.pathOptData + self.nameModel  + self.nameEnding + 'trialZ.npz'
        m.z       = np.load(dataread, allow_pickle=True)['sim_z']
        m.logs    = np.load(dataread, allow_pickle=True)['sim_log']
        print('files: ', np.load(dataread, allow_pickle=True).files)
        # m.edgePulse = np.load(dataread, allow_pickle=True)['edgePulse']
        return m
          
    def getPulses(self):
        dataIni        = self.pathOptData + self.nameModel + str(self.stages[0])  + self.nameEnding + '.npz'
        pulse = BaseData(dataIni, self.stages)
        pulse.d, pulse.ener = [[None]*pulse.nStg ]*pulse.nIsla, [[None]*pulse.nStg ]*pulse.nIsla
        for iXstg, stage in enumerate(self.stages):
            dataread        = self.pathOptData + self.nameModel + str(stage)  + self.nameEnding + '.npz'
            dataLog         = np.load(dataread, allow_pickle=True)['logsUDP']
            for iXisla, isla in enumerate(dataLog):
                pulse_ = np.array([log_[0] for log_ in isla])#.astype(np.float32)
                # pulse.ener[iXisla][iXstg]   = np.sum(np.abs(pulse_), axis=1)                                  
                pulse.d[iXisla][iXstg]      = pulse_
        return pulse
                
    def getFitness(self):
        dataIni        = self.pathOptData + self.nameModel + str(self.stages[0]) + self.nameEnding  + '.npz'
        fit = BaseData(dataIni, self.stages)
        fit.cost, fit.iXsort = [[None]*fit.nStg ]*fit.nIsla, [[None]*fit.nStg ]*fit.nIsla
        fit.raw, fit.renv, fit.sCor = [[None]*fit.nStg ]*fit.nIsla, [[None]*fit.nStg ]*fit.nIsla, [[None]*fit.nStg ]*fit.nIsla
        for iXstg, stage in enumerate(self.stages):
            dataread        = self.pathOptData + self.nameModel + str(stage)  + self.nameEnding + '.npz'
            dataLog         = np.load(dataread, allow_pickle=True)['logsUDP']
            for iXisla, isla in enumerate(dataLog):
                # fit.raw[iXisla][iXstg]   = np.array([log_[3] for log_ in isla]).astype(np.float32)
                fitCost_ = np.array([log_[5] for log_ in isla])#.astype(np.float32) 
                fit.cost[iXisla][iXstg]  = fitCost_
                # fit.iXsort[iXisla][iXstg]= np.argsort(fitCost_)
                fit.sCor[iXisla][iXstg]  = np.array([log_[6] for log_ in isla])
                fit.renv[iXisla][iXstg]  = np.array([log_[9] for log_ in isla])
        return fit
    
    def getPLI(self):
        dataIni        = self.pathOptData + self.nameModel + str(self.stages[0])  + self.nameEnding + '.npz'
        g = BaseData(dataIni, self.stages)
        pli = [[None]*g.nStg ]*g.nIsla
        for iXstg, stage in enumerate(self.stages):
            dataread        = self.pathOptData + self.nameModel + str(stage)  + self.nameEnding + '.npz'
            dataLog         = np.load(dataread, allow_pickle=True)['logsUDP']
            for iXisla, isla in enumerate(dataLog):
                # fit.raw[iXisla][iXstg]   = np.array([log_[3] for log_ in isla]).astype(np.float32)
                pli[iXisla][iXstg]  = np.array([log_[4] for log_ in isla])   
        return pli
        
    def getEnvelopes(self):
        dataIni        = self.pathOptData + self.nameModel + str(self.stages[0])  + self.nameEnding  + '.npz'
        env = BaseData(dataIni, self.stages)
        env.mStg = [[None]*env.nStg ]*env.nIsla
        env.mBu, env.relOn   = [[None]*env.nStg ]*env.nIsla, [[None]*env.nStg ]*env.nIsla
        for iXstg, stage in enumerate(self.stages):
            dataread        = self.pathOptData + self.nameModel + str(stage)  + self.nameEnding + '.npz'
            dataLog         = np.load(dataread, allow_pickle=True)['logsUDP']
            for iXisla, isla in enumerate(dataLog):
                env.mStg[iXisla][iXstg]   = np.array([log_[7] for log_ in isla]).astype(np.float32)
                env.mBu[iXisla][iXstg]    = np.array([log_[8] for log_ in isla]).astype(np.float32)
                env.relOn[iXisla][iXstg]  = np.array([log_[9] for log_ in isla]).astype(np.float32)       
        return env        
        
    def getOrders(self):
        dataIni        = self.pathOptData + self.nameModel + str(self.stages[0])  + self.nameEnding + '.npz'
        sync = BaseData(dataIni, self.stages)
        sync.g, sync.l = [[None]*sync.nStg ]*sync.nIsla, [[None]*sync.nStg ]*sync.nIsla
        sync.metaL, sync.metaG, sync.metaLm = [[None]*sync.nStg ]*sync.nIsla, [[None]*sync.nStg ]*sync.nIsla, [[None]*sync.nStg ]*sync.nIsla
        for iXstg, stage in enumerate(self.stages):
            dataread        = self.pathOptData + self.nameModel + str(stage)  + self.nameEnding + '.npz'
            dataLog         = np.load(dataread, allow_pickle=True)['logsUDP']
            for iXisla, isla in enumerate(dataLog):
                ord_    = np.array([log_[1] for log_ in isla]).astype(np.float32)
                sync.g[iXisla][iXstg]       = ord_[:,0]
                sync.l[iXisla][iXstg]       = ord_[:,1]
                sync.metaG[iXisla][iXstg]   = ord_[:,2] 
                metaL_ = np.array([log_[2] for log_ in isla]).astype(np.float32)
                sync.metaL[iXisla][iXstg]   = metaL_
                sync.metaLm[iXisla][iXstg]  = np.mean(metaL_, axis=1)
        return sync