# Copyright (C) 2016 EDF
# All Rights Reserved
# This code is published under the GNU Lesser General Public License (GNU LGPL)
import numpy as np
import StOptReg as reg

class TransitionStepRegressionDP:
    
    def __init__(self, p_pGridCurrent, p_pGridPrevious, p_pOptimize):
        
        self.m_pGridCurrent = p_pGridCurrent
        self.m_pGridPrevious = p_pGridPrevious
        self.m_pOptimize = p_pOptimize
        
    def oneStep(self, p_phiIn, p_condExp):
        
        nbRegimes = self.m_pOptimize.getNbRegime()
        phiOut = list(range(nbRegimes))
        nbControl = self.m_pOptimize.getNbControl()
        controlOut = list(range(nbControl))
        
        # only if the processor is working
        if self.m_pGridCurrent.getNbPoints() > 0:
            
            # allocate for solution
            for iReg in range(nbRegimes):
                phiOut[iReg] = np.zeros((p_condExp.getNbSimul(), self.m_pGridCurrent.getNbPoints()))
                
            for iCont in range(nbControl):
                controlOut[iCont] = np.zeros((p_condExp.getNbSimul(), self.m_pGridCurrent.getNbPoints()))
                
            # number of threads
            nbThreads = 1
            
            contVal = []
            
            for iReg in range(len(p_phiIn)):
                contVal.append(reg.ContinuationValue(self.m_pGridPrevious, p_condExp, p_phiIn[iReg]))
            
            # create iterator on current grid treated for processor
            iterGridPoint = self.m_pGridCurrent.getGridIteratorInc(0)
            
            # iterates on points of the grid
            for iIter in range(self.m_pGridCurrent.getNbPoints()):
                
                if iterGridPoint.isValid():
                    pointCoord = iterGridPoint.getCoordinate()
                    # optimize the current point and the set of regimes
                    solutionAndControl = self.m_pOptimize.stepOptimize(self.m_pGridPrevious, pointCoord, contVal, p_phiIn)
                    
                    # copy solution
                    for iReg in range(self.m_pOptimize.getNbRegime()):
                        phiOut[iReg][:,iterGridPoint.getCount()] = solutionAndControl[0][:,iReg]
                    
                    for iCont in range(nbControl):
                        controlOut[iCont][:,iterGridPoint.getCount()] = solutionAndControl[1][:,iCont]
                
                    iterGridPoint.nextInc(nbThreads)
        
        res = []
        res.append(phiOut)
        res.append(controlOut)
        return res
