#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Feb 23 09:51:31 2021

@author: Oscar Portoles
"""

import numpy as np
import sys
import visArchis_seq as v
# import optimProbs as oudp
import modelStages as udp
import copy

nSample = 1000
topNbest = 10000 #~the best 1%

d = v.LoadData()
d.nameModel = 'stgSequence_Lopt_allPu1500_wE_P040_optIni_v5_fixG0p15L0p6_Stg'
d.nameEnding =''
d.stages = [1,2,3,4] # to read file names
saveTo = './dataSave/relevanceOfPulses_P030_G0p15L0p6.npz'

# get data
pulse   = d.getPulses()
fit     = d.getFitness()
m       = d.getBestModel()
# Merge all islands
allPulse, allCost, allsCor, allRenv, allFC = [None]*pulse.nStg, [None]*pulse.nStg, [None]*pulse.nStg, [None]*pulse.nStg, [None]*pulse.nStg
for sIx in range(pulse.nStg):
    mergeP_, mergeC_, mergeScor_, mergeRenv_ = [], [], [], [] 
    for iIx in range(pulse.nIsla):
        mergeP_.append(pulse.d[iIx][sIx])
        mergeC_.append(fit.cost[iIx][sIx])
        mergeScor_.append(fit.sCor[iIx][sIx].ravel())
        mergeRenv_.append(np.squeeze(fit.renv[iIx][sIx]))
    allPulse[sIx] = np.vstack(mergeP_)
    allCost[sIx] = np.hstack(mergeC_)
    allsCor[sIx] = np.hstack(mergeScor_)
    allRenv[sIx] = np.vstack(mergeRenv_)

# Keep top N best solutions    
pulseTh, costTh, fcTh, sCorTh, rEnvTh = [None]*4, [None]*4, [None]*4, [None]*4, [None]*4
minCostTh = np.zeros(4)
for i in range(4):
    pulseTh[i]  = allPulse[i][ np.argsort(allCost[i])[-topNbest:], :]
    costTh[i]   = allCost[i][  np.argsort(allCost[i])[-topNbest:]]
    sCorTh[i]   = allsCor[i][  np.argsort(allCost[i])[-topNbest:]]
    rEnvTh[i]   = allRenv[i][  np.argsort(allCost[i])[-topNbest:], :]
    minCostTh[i]= np.min(costTh[i])
iXpick   = np.random.choice(np.arange(topNbest), size=nSample, replace=False)
  
del pulse, fit, allPulse, allCost, allRenv, allsCor

# best sequence of pulses and its scores
pulseBest = np.zeros((m.nroi,m.nStg))
fitBest   = np.zeros(m.nStg)
sCorBest  = np.zeros(m.nStg)
rEnvBest  = np.zeros((m.nStg,68))
for i in range(m.nStg):
    pulseBest[:,i]  = m.logs[i][0][0]
    fitBest[i]      = m.logs[i][0][3]
    sCorBest[i]     = m.logs[i][0][6] 
    rEnvBest[i,:]   = np.squeeze(m.logs[i][0][9])

fitRef  = np.zeros((nSample,4))
sCorRef = np.zeros((nSample,4))
rEnvRef = np.zeros((nSample,4,68))
pliRef  = np.zeros((nSample,2278,4))
fitTe   = np.zeros((nSample,68,4))
sCorTe  = np.zeros((nSample,68,4))
rEnvTe  = np.zeros((nSample,68,4,68))
pulseTe = np.zeros((nSample,68,4))

for iXt in range(nSample):
    for iXs in range(4):
        pulseRand = pulseTh[iXs][iXpick[iXt],:] # pulse to be tested
        pulseTe[iXt,:,iXs] = pulseRand
        stages = list(range(iXs+1))
        if len(stages) == 1:
            kao = udp.KAOmodel(stageIx=stages[0])
        else:
            kao = udp.KAOmodel(stageIx=stages)
        kao._setHistDelayAndContainers()
        kao._edgesImpulse()  
        kao._getEmpiReativEnv()
        kao.doNoise(variance=0)
        # do pulses
        pulseTest = copy.deepcopy(pulseBest[:,stages])
        pulseTest[:,iXs] = pulseRand
        kao.pulse = pulseTest
        
        kao.runModel_Pulsein_sequence() # noise is zeros
        
        fit_ = kao.similarityAllPLIs()
        fitRef[iXt, iXs] = fit_[iXs]
        pliRef[iXt, :, iXs] = kao.pli[iXs,:]
        
        kao.relativeEnvBuPre()
        kao.getScorrelRelEnv()
        sCorRef[iXt, iXs] = kao.sCor[iXs]   
        rEnvRef[iXt, iXs,:] = kao.relEnv[:,iXs]  
        # iterate region to get relevance with a missing pulse
        for iXroi in range(68):
            if len(stages) == 1:
                kao = udp.KAOmodel(stageIx=stages[0])
            else:
                kao = udp.KAOmodel(stageIx=stages)            
            kao._setHistDelayAndContainers()
            kao._edgesImpulse()  
            kao._getEmpiReativEnv()
            kao.doNoise(variance=0)
            # do pulses
            pulseTest = copy.deepcopy(pulseBest[:,stages])
            pulseTest[:,iXs] = pulseRand
            pulseTest[iXroi,iXs] = 0
            kao.pulse = pulseTest
            
            kao.runModel_Pulsein_noise()
    
            fit_ = kao.similarityAllPLIs()
            fitTe[iXt, iXroi, iXs] = fit_[iXs]
            
            kao.relativeEnvBuPre()
            kao.getScorrelRelEnv()
            sCorTe[iXt, iXroi, iXs]   = kao.sCor[iXs]
            rEnvTe[iXt, iXroi, iXs,:] = kao.relEnv[:,iXs]
            
np.savez(saveTo, rEnvTe=rEnvTe, sCorTe=sCorTe, fitTe=fitTe, pulseTe=pulseTe, pulseTh=pulseTh,costTh=costTh,iXpick=iXpick,
         fitRef=fitRef, sCorRef=sCorRef, pliRef=pliRef, rEnvRef=rEnvRef, sCorTh=sCorTh, rEnvTh=rEnvTh,
         pulseBest=pulseBest, fitBest=fitBest, sCorBest=sCorBest, rEnvBest=rEnvBest)