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

@author: Oscar Portoles
"""

import numpy as np
from modelStages import WBLSmodel
import copy

    
class Opt_allROIpulse_Wenecor(WBLSmodel):
    """OPtimization of pulses for several stages sequetially
        With or without time delays (looks at self.velocity)
        OPtimizes FC, energy and bump topology"""
    def __init__(self,stageIx=0,vel=0,zh=None):
        super().__init__(stageIx=stageIx,vel=vel,zh=zh)
        
        self.widthPulse   = 0.040   # [seconds] distance from center of bumps/pulse to state edges, width of pulse
        self.offMidBump   = 0.01    # relative to the middel poit of a bump
        
        self._setHistDelayAndContainers()
        self._edgesImpulse()  
        self._getEmpiReativEnv()
        
    def get_name(self):
        return "Opt. all ROI pulses, constrains, no delays, fixed G,L"
    
    def get_bounds(self):
        """Boundaries on: pulses """
        self.lowbound   = [-1000]*self.C
        self.upbound    = [ 1000]*self.C
        self.maxEner    = np.max([self.upbound[0], np.abs(self.lowbound[0])]) * self.C
        return (self.lowbound, self.upbound)
          
    def fitness(self, x):
        self.pulse      = np.float64(x)
        if self.velocity == 0:
            self.runModelNoDel_Pulsein()
        else:
            self.runModel_Pulsein()     
        # cost function
        self.relativeEnvBuPre()
        self.getScorrelRelEnv()           
        fitByStage = self.similarityPLIsWenergySperm()
        # for constrains
        self.meanEnvBuStage()
        fitSum = np.sum(fitByStage)
        if np.isnan(fitSum) or ( np.median(self.msEnv) < 0.01 ).any() or self.sCor < 0: # Penalize wrong solutions
            fitSum = -1
        # Keep logs
        self.log.append( [np.float64(x),self.orders, self.metaL,fitByStage, self.pli.astype(np.float16), np.float16(fitSum), np.float16(self.sCor),
                          self.msEnv.astype(np.float16), self.mbEnv.astype(np.float16), self.relEnv.astype(np.float16)] )
        return np.array([-1* fitSum])

          
class FindAndKeepBestSolution(WBLSmodel):
    """ Extracts the logs form an optimization, get parameters from best solution, 
        run model and keep last samples. Used for squence of stages"""
        
    def extractBestPulse(self, logs):
        maxFit = 0
        for logIsla in logs:
            fit = np.array([log_[3] for log_ in logIsla]).astype(np.float32)
            if np.max(fit) > maxFit:
                iXmax = np.argmax(fit)
                self.bestPulse = np.array([log_[0] for log_ in logIsla])[iXmax,:].astype(np.float64)
                
    def getBestWeightedModel(self, logs, per):
        """ @ logs: list, logs from UDP
            @ per: float, percentile of best fitness to look for best solution"""
        fit, corr, ener, pulse = [], [], [], []
        for logIsla in logs:
            fit.append(  np.array([log_[5] for log_ in logIsla]).ravel() )
            ener.append( np.array([np.sum(np.abs(log_[0])) for log_ in logIsla] ).ravel() )
            corr.append( np.array([log_[6] for log_ in logIsla]).ravel() )
            pulse.append(np.array([log_[0] for log_ in logIsla] ) )
            
        fit     = np.hstack(fit)
        ener    = np.hstack(ener)
        corr    = np.hstack(corr)
        pulse   = np.vstack(pulse) 
        # top fitness pertcentile threshold
        perTh   = np.percentile(fit, per)
        mask    = fit >= perTh
        fit, corr, ener, pulse = fit[mask], np.float32(corr[mask]), np.float32(ener[mask]), pulse[mask,:]
        # best correlation-energy trade-off
        iXbest = np.argmin( ener*(1-corr))
        self.bestPulse  = pulse[iXbest, :].astype(np.float64)
        self.bestEnergy, self.bestFit, self.bestCorr = ener[iXbest], fit[iXbest], corr[iXbest]
        
    def runBestModel(self):
        self._setHistDelayAndContainers()
        self._edgesImpulse() 
        self._getEmpiReativEnv()
        self.pulse      = copy.deepcopy(self.bestPulse)
        self.maxEner    = 1000.0 * self.C
        if self.velocity == 0:
            self.runModelNoDel_Pulsein()
        else:
            self.runModel_Pulsein()
        # cost function
        self.relativeEnvBuPre()
        self.getScorrelRelEnv()           
        fitByStage = self.similarityPLIsWenergySperm()
        # for constrains
        self.meanEnvBuStage()
        fitSum = np.sum(fitByStage)
        # Keep logs
        self.log.append( [self.pulse,self.orders, self.metaL,fitByStage, self.pli.astype(np.float16), np.float16(fitSum), np.float16(self.sCor),
                          self.msEnv.astype(np.float16), self.mbEnv.astype(np.float16), self.relEnv.astype(np.float16)] )
        return self.z[:, -(self.maxDlay+1):, :]