import datetime
import os
from logging import debug, info

import numpy as np
import pandas as pd
import xarray as xr
from netCDF4 import Dataset

from ....utils.dates import date_range
from ....utils.netcdf import readnc



[docs]
def ini_periods(self, **kwargs):
    datei = self.datei
    datef = self.datef
    nho = self.nho
    nhours = self.nhours

    # List of sub-simulation windows
    self.subsimu_dates = date_range(datei, datef, period=self.periods)

    # Check that subperiods are all of the same length
    if np.unique(np.diff(self.subsimu_dates)).size > 1:
        raise Exception(
            "Trying to run CHIMERE on a simulation window not fitting the "
            "sub-simulation length. \nPlease check the compatibility between "
            "the duration of your sub-simulations ({}) and the simulation "
            "window ({} to {}): \n"
            "Guessed sub-periods: \n".format(self.periods, datei, datef)
            + "\n".join(
                [
                    "  - {}: {}".format(d0, d1 - d0)
                    for d0, d1 in zip(self.subsimu_dates[:-1], self.subsimu_dates[1:])
                ]
            )
        )

    # Time steps defined by the user
    nphour_ref = self.nphour_ref

    # Numbers of time steps per hour really used in the simulation
    self.tstep_dates = {}
    self.tstep_all = []
    self.nhour = {}
    self.subtstep = {}
    self.input_dates = {}

    # Stop-or-more function
    list_to_stop = self.stopORmore
    list_spec_stop = []
    list_thld_stop = []
    for l in list_to_stop:
        list_spec_stop.append(l[0])
        list_thld_stop.append(l[1])

    if list_thld_stop != [] and list_spec_stop != []:
        info("Looking at concentrations small enough to stop simulations")

    self.list_spec_stop = list_spec_stop
    self.list_thld_stop = list_thld_stop

    # Loop over sub simulations
    for dd in self.subsimu_dates[:-1]:
        # time-steps in METEO.nc, computed by diagmet
        try:
            met = dd.strftime(f"{self.meteo.dir}/{self.meteo.file}")
            debug(f"Reading {met} to initialize nphourm")
            ds_met = xr.open_dataset(met)
            if "nphourm2" in ds_met.variables:
                nbstep = ds_met.variables["nphourm2"].values.astype(int)
            elif "nphourm" in ds_met.variables:
                nbstep = ds_met.variables["nphourm"].values.astype(int)
            else:
                raise Exception(
                    f"The file {met} is available, but no information is "
                    f"available about nphourm. Please check your file"
                )
        except IOError:
            debug(f"{met} is not available. Using nphourref (={nphour_ref}).")
            nbstep = nhours * [nphour_ref]
        except AttributeError:
            debug(
                "No meteo plugin was defined in the Yaml. "
                f"Using default plugin with nphourref (={nphour_ref})."
            )
            nbstep = nhours * [nphour_ref]

        # Loop on hours and check CFL
        self.tstep_dates[dd] = []
        self.nhour[dd] = []
        self.subtstep[dd] = []
        for nh in range(nhours):
            nphour = nbstep[nh] if nphour_ref < nbstep[nh] else nphour_ref

            # Frequency in seconds
            # TODO: Check with FORTRAN: nphour rounding?
            freq = "{}s".format(int(3600 // nphour))

            # List of time steps
            # TODO: what about chemical time steps?
            # the time step to really use is nphour*ichemstep
            ddhi = dd + datetime.timedelta(hours=nh)
            ddhe = dd + datetime.timedelta(hours=nh + 1)
            drange = list(pd.date_range(ddhi, ddhe, freq=freq).to_pydatetime())
            self.tstep_dates[dd].extend(drange[:-1])
            self.tstep_all.extend(drange[:-1])

            # Saving substep indexes for matching with observation
            nphour_int = len(drange) - 1
            self.subtstep[dd].extend(list(range(1, nphour_int + 1)))
            self.nhour[dd].extend(nphour_int * [nh + 1])

        # List of dates for which inputs are needed
        self.input_dates[dd] = pd.date_range(
            dd, periods=nhours + 1, freq="1h"
        ).to_pydatetime()

        # Include last time step
        self.tstep_dates[dd].append(ddhe)
        self.tstep_dates[dd] = np.array(self.tstep_dates[dd])

    # Include very last time step
    self.tstep_all.append(ddhe)
    self.tstep_all = np.array(self.tstep_all)

    # Initializes dictionary to keep in memory whether observations were
    # already dumped for a given period
    self.iniobs = {ddi: False for ddi in self.subsimu_dates}
    self.nbobs_prior = {ddi: 0 for ddi in self.subsimu_dates}
    self.nbdatatot_prior = {ddi: 0 for ddi in self.subsimu_dates}
    self.reset_obs = {ddi: True for ddi in self.subsimu_dates}

    # Keep track for stop-or-more function (initialize run = True)
    self.runsimu = {ddi: True for ddi in self.subsimu_dates}

    # Keep in memory whether a given period has a successor period
    # The info is used by the adjoint to fetch or not the aend file
    self.chain = {ddi: True for ddi in self.subsimu_dates[:-1]}
    self.chain[self.subsimu_dates[-2]] = False