# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/00_core.ipynb (unless otherwise specified).

__all__ = ['AquaCropModel', 'perform_timestep']

# Cell
import matplotlib.pyplot as plt
import numpy as np
import pathlib
import os
import pandas as pd
import sys
[sys.path.append(i) for i in ['.', '..']]

# Cell
from .initialize import *
from .timestep import *
from .classes import *

# Cell
class AquaCropModel:

    def __init__(self,SimStartTime,SimEndTime,wdf,Soil,Crop,
                     Farmer=None,FieldMngt=None,FallowFieldMngt=None,
                     GwStruct=None,InitWC=None,planting_dates=None,
                     harvest_dates=None):

        self.SimStartTime = SimStartTime
        self.SimEndTime = SimEndTime
        self.wdf = wdf
        self.Soil = Soil
        self.Crop = Crop
        self.planting_dates = planting_dates
        self.harvest_dates = harvest_dates

        self.Farmer = Farmer
        self.FieldMngt = FieldMngt
        self.FallowFieldMngt = FallowFieldMngt
        self.GwStruct = GwStruct
        self.InitWC = InitWC

        if Farmer == None:  self.Farmer = FarmerClass(IrrMethod=0);
        if FieldMngt == None:  self.FieldMngt = FieldMngtClass();
        if FallowFieldMngt == None:  self.FallowFieldMngt = FieldMngtClass();
        if GwStruct == None:  self.GwStruct = GwClass();
        if InitWC == None:  self.InitWC = InitWCClass();



    def initialize(self,):
        """
        Initialize variables

        Looks for filepaths within aquacrop/data/`fileDir`/Input/FileSetup.txt
        and stores inside a `FileLocationsClass`

        """

        # define model runtime
        self.ClockStruct = read_clock_paramaters(self.SimStartTime,self.SimEndTime)

        # get weather data
        self.weather_df = read_weather_inputs(self.ClockStruct,self.wdf)

        # read model params
        self.ClockStruct, self.ParamStruct = read_model_parameters(self.ClockStruct,self.Soil,
                                                                   self.Crop,self.planting_dates,self.harvest_dates)

        # read irrigation management
        self.ParamStruct,self.Farmer = read_irrigation_management(self.ParamStruct,self.Farmer,self.ClockStruct)

        # read field management
        self.ParamStruct = read_field_management(self.ParamStruct,self.FieldMngt,self.FallowFieldMngt)

        # read groundwater table
        self.ParamStruct = read_groundwater_table(self.ParamStruct,self.GwStruct,self.ClockStruct)

        # Compute additional variables
        self.ParamStruct = compute_variables(self.ParamStruct,self.weather_df,self.ClockStruct)

        # read, calculate inital conditions
        self.ParamStruct, self.InitCond = read_model_initial_conditions(self.ParamStruct,self.ClockStruct,self.InitWC)

        self.ParamStruct = create_soil_profile(self.ParamStruct)

        self.InitCond.ParamStruct = self.ParamStruct

        Outputs = OutputClass()
        Outputs.Water = np.zeros((len(self.ClockStruct.TimeSpan),15))
        Outputs.Flux = np.zeros((len(self.ClockStruct.TimeSpan),16))
        Outputs.Growth = np.zeros((len(self.ClockStruct.TimeSpan),13))
        Outputs.Final = pd.DataFrame(columns = ['Season','Name','HarvestDate','Step','Yield','Tirr'])


        self.Outputs=Outputs

        # save model weather to InitCond
        self.InitCond.weather = self.weather_df.values



        #return self.ClockStruct,self.InitCond,self.Outputs
        return


    def step(self,num_steps=1,till_termination=False):

        if till_termination==True:

            while self.ClockStruct.ModelTermination == False:

                self.ClockStruct,self.InitCond,self.Outputs = perform_timestep(self.ClockStruct,
                                                                               self.InitCond,
                                                                               self.Outputs,
                                                                               self.Farmer)
        else:

            for i in range(num_steps):

                self.ClockStruct,self.InitCond,self.Outputs = perform_timestep(self.ClockStruct,
                                                                               self.InitCond,
                                                                               self.Outputs,
                                                                               self.Farmer)

                if self.ClockStruct.ModelTermination: return

        #return self.ClockStruct,self.InitCond,self.Outputs
        return

# Cell
def perform_timestep(ClockStruct,InitCond,Outputs,Farmer):

    """
    Function to run a single time-step (day) calculation of AquaCrop-OS

    *Arguments:*\n


    `ClockStruct` : `ClockStructClass` :  model time paramaters

    `weather_df`: `pandas.DataFrame` :  weather data for simulation period

    `InitCond` : `InitCondClass` :  containing all initiall model paramaters

    `Outputs` : `OutputClass` :  object to store outputs


    *Returns:*

    `ClockStruct` : `ClockStructClass` :  updated time paramaters

    `InitCond` : `InitCondClass` :  updated model paramaters

    `Outputs` : `OutputClass` :  updated with the days outputs



    """


    # extract weather data for current timestep
    #weather_step = weather_df[weather_df.Date==ClockStruct.StepStartTime]
    weather_step = InitCond.weather[ClockStruct.TimeStepCounter]

    #%% Get model solution %%
    NewCond,Outputs = solution(InitCond,ClockStruct,weather_step,Outputs,Farmer)

    #%% Check model termination %%
    ClockStruct = check_model_termination(ClockStruct,NewCond)

    #%% Update time step %%
    ClockStruct, InitCond,Outputs = update_time(ClockStruct,NewCond,Outputs)

    return ClockStruct,InitCond,Outputs