import ast
import itertools
import math

import arrow as ar
import numpy as np
import pandas as pd
from heapq import nsmallest
from pathlib import Path
from tqdm import tqdm

from geoprepare import base
from geocif.agmet import plot, utils


class AgmetGeo(base.BaseGeo):
    """Lightweight geo object for agmet plots. No Neptune dependency."""

    def __init__(self, path_config_file):
        super().__init__(path_config_file)
        self.parse_config()

    def _get_option(self, option, default=None, sections=("AGMET", "DEFAULT")):
        """Get a config option, checking sections in order."""
        for section in sections:
            if self.parser.has_option(section, option):
                return ast.literal_eval(self.parser.get(section, option))
        return default

    def parse_config(self, section="DEFAULT"):
        self.project_name = self.parser.get("PROJECT", "project_name")
        super().parse_config(project_name=self.project_name, section="DEFAULT")

        self.countries = ast.literal_eval(self.parser.get("DEFAULT", "countries"))
        self.forecast_seasons = self._get_option(
            "forecast_seasons", [ar.utcnow().year]
        )
        self.plot_seasons = self._get_option(
            "plot_seasons", self.forecast_seasons
        )
        self.models = self._get_option("models", [])
        self.eo_plot = self._get_option(
            "eo_plot",
            ["ndvi", "cpc_tmax", "cpc_tmin", "chirps", "esi_4wk", "soil_moisture_as1", "soil_moisture_as2"],
        )
        self.eo_model = self._get_option(
            "eo_model",
            ["ndvi", "cpc_tmax", "cpc_tmin", "chirps", "esi_4wk", "soil_moisture_as1", "soil_moisture_as2"],
        )

        # Logos (optional)
        try:
            self.logo_harvest = self.dir_metadata / self.parser.get(
                "METADATA", "logo_harvest"
            )
            self.logo_geoglam = self.dir_metadata / self.parser.get(
                "METADATA", "logo_geoglam"
            )
        except Exception:
            self.logo_harvest = None
            self.logo_geoglam = None

    def read_statistics(self, country=None, read_countries=False, **kwargs):
        """Read zone/country info. country is optional when only reading countries."""
        if read_countries:
            path_countries = (
                self.dir_input / "statistics" / self.parser.get("DEFAULT", "zone_file")
            )
            self.df_countries = (
                pd.read_csv(path_countries)
                if path_countries.is_file()
                else pd.DataFrame()
            )
        if country is not None:
            super().read_statistics(country, read_countries=read_countries, **kwargs)

    def get_calendar_region_for_region(self, df, region):
        return df[df["region"] == region]["calendar_region"].values[0]

    def setup_country(self, country, scale, crop, growing_season):
        self.country = country
        self.scale = scale
        self.crop = crop
        self.growing_season = growing_season
        self.category = self.parser.get(country, "category")
        self.use_cropland_mask = self.parser.getboolean(country, "USE_CROPLAND_MASK")

        self.get_dirname(country)
        self.get_ccs_dataframe(country, scale, crop, growing_season)

        self.list_regions = self.df_ccs["region"].unique()
        self.list_calendar_regions = self.df_ccs["calendar_region"].unique()

        self.precip_var = (
            "chirps" if "chirps" in self.df_ccs.columns.values else "cpc_precip"
        )

        self.dir_output = self.dir_output / ar.now().format("MMMM_DD_YYYY")

    def setup_region(self, region, plot_season, type_region="region"):
        self.region = region
        self.calendar_region = self.get_calendar_region_for_region(self.df_ccs, region)
        self.plot_season = plot_season
        self.type_region = type_region

        if type_region == "region":
            self.df_region = self.df_ccs[self.df_ccs["region"] == region]
        elif type_region == "calendar_region":
            self.df_region = self.df_ccs[
                self.df_ccs["calendar_region"] == self.calendar_region
            ]
        elif type_region == "region_year":
            self.df_region = self.df_ccs[
                (self.df_ccs["region"] == region) & (self.df_ccs["year"] == plot_season)
            ]
        elif type_region == "calendar_region_year":
            self.df_region = self.df_ccs[
                (self.df_ccs["calendar_region"] == self.calendar_region)
                & (self.df_ccs["year"] == plot_season)
            ]
        else:
            raise ValueError(f"Unknown type_region: {type_region}")

        self.df_region.index = pd.to_datetime(self.df_region.index)

        folder = f"{self.crop}_s{self.growing_season}_{self.plot_season}"
        self.dir_agmet = (
            self.dir_output
            / "crop_condition"
            / self.category
            / self.country
            / folder
            / "condition"
        )

        (
            self.date_planting,
            self.date_greenup,
            self.date_senescence,
            self.date_harvesting,
        ) = self.get_calendar(self.region, self.plot_season)

    def create_run_combinations(self):
        all_combinations = []

        for country in self.countries:
            scales = ast.literal_eval(self.parser.get(country, "scales"))
            crops = ast.literal_eval(self.parser.get(country, "crops"))
            growing_seasons = ast.literal_eval(
                self.parser.get(country, "growing_seasons")
            )

            for scale in scales:
                for crop in crops:
                    for growing_season in growing_seasons:
                        all_combinations.extend(
                            list(
                                itertools.product(
                                    [country], [scale], [crop], [growing_season]
                                )
                            )
                        )

        return all_combinations

    def get_calendar(self, region, forecast_season):
        SEASON = "harvest_season"
        CAL = "crop_calendar"

        df_sub = self.df_ccs[self.df_ccs["region"] == region]
        df_sub.index = pd.to_datetime(df_sub.index)
        sr_cal = df_sub[df_sub[SEASON] == forecast_season][["doy", CAL]]

        sr_cal[CAL] = pd.to_numeric(sr_cal[CAL], errors="coerce")

        if sr_cal.empty:
            return np.NaN, np.NaN, np.NaN, np.NaN
        else:
            date_planting = (sr_cal[CAL] == 1).idxmax()
            date_greenup = (
                (sr_cal[CAL] == 2).idxmax() if len(sr_cal[sr_cal[CAL] == 2]) else None
            )
            date_senesc = (
                (sr_cal[CAL][::-1] == 2).idxmax()
                if len(sr_cal[sr_cal[CAL] == 3])
                else None
            )
            date_harvesting = (
                (sr_cal[CAL][::-1] == 3).idxmax()
                if len(sr_cal[sr_cal[CAL] == 3])
                else None
            )

            return date_planting, date_greenup, date_senesc, date_harvesting

    def get_ccs_dataframe(self, country, scale, crop, growing_season):
        dir_ccs = self.dir_output / self.dir_threshold / country

        self.df_ccs = pd.read_csv(
            dir_ccs / f"{country}_{crop}_s{growing_season}.csv", index_col=0
        )

        self.df_ccs["datetime"] = pd.to_datetime(self.df_ccs.index)
        self.df_ccs.index.name = None

    def get_closest_season(self, season):
        self.closest = nsmallest(
            5 + 1, range(2001, ar.utcnow().year), key=lambda x: abs(x - season)
        )

        if season in self.closest:
            self.closest.remove(season)

    def check_date(self, df, plot_season):
        last_valid_date = pd.to_datetime(df["chirps"].last_valid_index()).date()

        bool_year_check = ar.utcnow().year <= plot_season
        bool_date_check = (last_valid_date > self.date_planting.date()) & (
            last_valid_date < self.date_harvesting.date()
        )

        return bool_year_check, bool_date_check

    def add_precip_forecast(self, plot_season):
        bool_year_check, bool_date_check = self.check_date(self.df_ccs, plot_season)

        if bool_year_check & bool_date_check:
            base_dir = self.dir_output / self.dir_threshold / self.country / self.scale
            path_gefs = (
                base_dir / "cr" / "chirps_gefs"
                if self.use_cropland_mask
                else base_dir / self.crop / "chirps_gefs"
            )

            df_gefs = pd.read_csv(
                list(path_gefs.glob(self.df_region["region"].unique()[0] + "*.csv"))[0],
                header=None,
            )
            val_gefs = (
                float(df_gefs.values[1][5])
                if not math.isnan(float(df_gefs.values[1][5]))
                else 0.0
            )

            start_date = ar.utcnow().shift(days=+15).date()
            end_date = ar.utcnow().shift(days=+16).date()

            self.df_region.loc[start_date:end_date, "chirps_gefs"] = val_gefs
            self.df_ccs.loc[start_date.strftime("%Y-%m-%d"), "chirps_gefs"] = val_gefs


def create_title_for_plot(obj):
    """
    Args:
        obj:

    Returns:

    """
    region_name = obj.region.replace("_", " ").title()
    calendar_region_name = obj.calendar_region.replace("_", " ").title()
    country_name = obj.country.replace("_", " ").title()
    long_crop_name = utils.get_crop_name(obj.crop)

    # Get name of crop based on metadata/crop_per_season.csv file
    df_crop_per_season = pd.read_csv(obj.dir_metadata / "crop_per_season.csv")

    crop_name = df_crop_per_season[
        (df_crop_per_season["country"] == obj.country)
        & (df_crop_per_season["crop"] == obj.crop)
        & (df_crop_per_season["season"] == int(obj.growing_season))
    ]["name"].values

    crop_name = crop_name[0] if crop_name else long_crop_name.replace("_", " ").title()

    title_line_1 = f"{region_name} ({calendar_region_name}, {country_name})"
    title_line_2 = f"{crop_name} {obj.plot_season}"

    sup_title = f"{title_line_1}\n{title_line_2}"

    return sup_title


def loop_agmet(path_config_file=None):
    """
    Args:
    """
    # Create geo object
    obj = AgmetGeo(path_config_file)

    # Read in data on crop names in each country
    obj.read_statistics(read_countries=True)

    # Create combinations of run parameters
    all_combinations = obj.create_run_combinations()

    pbar = tqdm(all_combinations, total=len(all_combinations))
    for country, scale, crop, growing_season in pbar:
        # Setup country information
        obj.setup_country(country, scale, crop, growing_season)

        # Loop through seasons and plot for each admin_1 region and calendar region
        for plot_season in obj.plot_seasons:
            # Get list of the years that are closest to the year which we want to plot
            obj.get_closest_season(plot_season)

            for region in obj.list_regions:
                # Setup the region information
                obj.setup_region(region, plot_season, "region")

                # If available, add CHIRPS-GEFS data to the dataframe
                if obj.precip_var == "chirps":
                    obj.add_precip_forecast(plot_season)

                # Get crop calendar dates for region
                dates_calendar = [
                    obj.date_planting,
                    obj.date_greenup,
                    obj.date_senescence,
                    obj.date_harvesting,
                ]

                # TODO: Only plot if dates_calendar has valid dates
                # Create title for plot
                sup_title = create_title_for_plot(obj)

                ###############################################################
                # Agmet plots for regions
                ###############################################################
                plot.plots_ts_cur_yr(
                    obj.df_region,
                    obj.eo_plot,
                    closest=obj.closest,
                    dates_cal=dates_calendar,
                    frcast_yr=obj.plot_season,
                    logos=[obj.logo_harvest, obj.logo_geoglam],
                    dir_out=obj.dir_agmet / obj.scale,
                    sup_title=sup_title,
                    fname=f"{obj.region}.png",
                )

                ###############################################################
                # Agmet plots for calendar regions
                ###############################################################
                columns = obj.eo_model + ["month", "day", "yield"]
                if "chirps" in obj.df_region.columns:
                    bool_year_check, bool_date_check = obj.check_date(
                        obj.df_region, obj.plot_season
                    )

                    if bool_date_check and bool_year_check:
                        columns = obj.eo_model + [
                            "month",
                            "day",
                            "chirps_gefs",
                            "yield",
                        ]

                df_calendar_region = (
                    obj.df_region.groupby(
                        [
                            "country",
                            "calendar_region",
                            "harvest_season",
                            "doy",
                            "datetime",
                        ]
                    )[columns]
                    .mean()
                    .reset_index()
                )
                df_calendar_region.loc[:, "average_temperature"] = (
                    df_calendar_region["cpc_tmax"] + df_calendar_region["cpc_tmin"]
                ) / 2.0
                df_calendar_region.set_index(
                    pd.DatetimeIndex(df_calendar_region["datetime"]),
                    inplace=True,
                    drop=True,
                )
                df_calendar_region.index.name = None
                df_calendar_region.sort_values(by="datetime", inplace=True)

                # TODO: Only plot if dates_calendar has valid dates
                # Create title for plot
                sup_title = create_title_for_plot(obj)

                if not df_calendar_region.empty:
                    plot.plots_ts_cur_yr(
                        df_calendar_region,
                        obj.eo_plot,
                        closest=obj.closest,
                        dates_cal=dates_calendar,
                        frcast_yr=obj.plot_season,
                        logos=[obj.logo_harvest, obj.logo_geoglam],
                        dir_out=obj.dir_agmet / "district",
                        sup_title=sup_title,
                        fname=f"{obj.calendar_region}.png",
                    )


def run(path_config_files=[]):
    loop_agmet(path_config_files)


if __name__ == "__main__":
    run()
