#!/usr/bin/python3

# 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
import StOptGrids
import StOptGlobal
import dp.DynamicProgrammingByRegression as dyn
import unittest
import simulators.AR1Simulator as ar1
import dp.OptimizeLake as ol
import dp.SimulateRegressionControlUsingControl as simpart



def almost_equal(x, y, eps=1e-5):
    return abs(x - y) <= eps

class ZeroFunction:

    def __init__(self):

        return None

    def set(self, a, b, c):

        return 0.

# valorization of a given Lake on a  grid
# Gain are proportional to what is withdrawed from the storage
# Only inflows are stochastic
# p_grid             the grid
# p_maxLevelStorage  maximum level
# p_mesh             number of mesh
# p_bCheckClose      Do we check if optimization and simulations are close
def lake(p_grid, p_maxLevelStorage, p_mesh, p_bCheckClose):

    # storage
    #########
    withdrawalRateStorage = 1000

    maturity = 1.
    nstep = 10

    # number of simulations
    nbsimulOpt = 8000
    nbsimulSim = 8000

    # inflow model
    D0 = 50 # initial inflow
    m = D0 # average inflow
    sig = 5. # volatility
    mr = 5. # mean reverting
    # a backward simulator
    ######################
    backSimulator = ar1.AR1Simulator(D0, m, sig, mr, maturity, nstep, nbsimulOpt, False)
    # optimizer
    ##########
    storage = ol.OptimizeLake(withdrawalRateStorage)

    # regressor
    ###########
    nbMesh = np.array([], dtype = np.int32)

    if p_mesh > 0:
        nbMesh = np.zeros(1, dtype = np.int32) + p_mesh

    regressor = reg.LocalLinearRegression(nbMesh)
    # final value
    vFunction = ZeroFunction()

    # initial values
    initialStock = np.zeros(1) + p_maxLevelStorage
    initialRegime = 0 # only one regime

    # Optimize
    fileToDump = "CondExpLake"

    # link the simulations to the optimizer
    storage.setSimulator(backSimulator)
    valueOptim = dyn.DynamicProgrammingByRegression(p_grid, storage, regressor , vFunction, initialStock, initialRegime, fileToDump)

    # Forward simulation using the control :
    forSimulator = ar1.AR1Simulator(D0, m, sig, mr, maturity, nstep, nbsimulSim, True)
    storage.setSimulator(forSimulator)
    valueSimulation = simpart.SimulateRegressionControlUsingControl(p_grid, storage, vFunction, initialStock, initialRegime, fileToDump)

    return valueOptim, valueSimulation


class testLakeTest(unittest.TestCase):

    def setUp(self):
        np.random.seed(0)

    # linear interpolation
    def test_simpleStorageLegendreLinear(self):

        # storage
        #########
        maxLevelStorage = 5000
        # grid
        ######
        nGrid = 10
        lowValues = np.zeros(1)
        step = np.zeros(1) + (maxLevelStorage / nGrid)
        nbStep = np.zeros(1, dtype = np.int32) + nGrid
        poly = np.zeros(1, dtype = np.int32) + 1
        grid = StOptGrids.RegularLegendreGrid(lowValues, step, nbStep, poly)
        nbmesh = 4
        accuracyEqual = 100
        optim, simulationUsingControl  = lake(grid, maxLevelStorage, nbmesh, True)

        self.assertAlmostEqual(optim , simulationUsingControl, None, None, accuracyEqual)


if __name__ == '__main__':
    unittest.main()