# Copyright (C) 2016 EDF
# All Rights Reserved
# This code is published under the GNU Lesser General Public License (GNU LGPL)
import StOptSDDP
import StOptSDDPUnitTest
import numpy as NP
import unittest
import StOptGeners

# unitest for SDDP framework
############################


class FakeOptimizer(unittest.TestCase):

    def test(self):

        # Create a fake optimizer depending of a fake simulator
        simul =StOptSDDPUnitTest.SimulTest()
        optim= StOptSDDPUnitTest.OptimizeTest(simul)

        # final cuts
        finCut =StOptSDDP.SDDPFinalCut()

        # SDDP problem parameters
        nbSimulCheckForSimu = 1
        initialState = NP.array([1.,2.])
        dates  = NP.array([1.,2.])
        nbMesh= NP.array([1,1],NP.int32)
        nameRegressor =" "
        nameCut = " "
        nameVisitedStates = " "
        iter= 1
        accuracy=0.01
        nStepConv=10

        # test backward and forward mapping
        sol = StOptSDDPUnitTest.backwardForwardSDDPTestMapping(optim,nbSimulCheckForSimu,initialState,finCut,dates,nbMesh,nameRegressor,nameCut,nameVisitedStates,iter,accuracy,nStepConv)


class LocalLinearRegressionForSDDP(unittest.TestCase):

    def test(self):

        # Create a local linear regressor for SDDP
        particles= NP.array([[1.,2.],[3.,3.]],NP.float)
        mesh = NP.array([1,2],NP.int32)
                           
        regressor =StOptSDDP.LocalLinearRegressionForSDDP(1,particles,mesh)
        
        archiveToWrite = StOptGeners.BinaryFileArchive("MyArchive","w")

        # dump SDDP regressor
        archiveToWrite.dump(regressor)

        # visited state
        aVisitedState = StOptSDDP.SDDPVisitedStates()
        # a state
        aState =  NP.array([3.,4.],NP.float)
        # a particle
        aParticle = NP.array([1.,2.],NP.float)
        aVisitedState.addVisitedState(aState, aParticle,regressor)

        # dump visited state
        archiveToWrite.dump(aVisitedState)

        # test backward SDDP mapping
        simul =StOptSDDPUnitTest.SimulTest()
        optim= StOptSDDPUnitTest.OptimizeTest(simul)
        # get backward base simulator
        simBase = optim.getSimulatorBackward()
        dates = NP.array([1.],NP.float)
        initialState=  NP.array([1.],NP.float)
        finalCut =StOptSDDP.SDDPFinalCut()
        archiveRegresssor = StOptGeners.BinaryFileArchive("MyArchiveReg","w")
        archiveCut  = StOptGeners.BinaryFileArchive("MyArchiveCut","r")
        a = StOptSDDPUnitTest.backwardSDDPTestMapping(optim,simBase,dates,initialState,finalCut,archiveRegresssor,"MyString",archiveCut)

        # test backward mapping
        b= StOptSDDPUnitTest.forwardSDDPTestMapping(optim,simBase,dates,initialState,finalCut,1,archiveRegresssor,archiveCut,"MyString")
        
        
if __name__ == '__main__': 
    unittest.main() 
