import ast

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

from geoprepare import base
from geocif.agmet import plot, utils
from geocif.ml import stats
from geocif import utils as ut

# Show usage info on import
from rich.console import Console
from rich.panel import Panel
from rich.table import Table

_console = Console()
_table = Table(show_header=False, box=None, padding=(0, 1))
_table.add_column(style="bold cyan", no_wrap=True)
_table.add_column()
_table.add_row("Usage", "from geocif.agmet import geoagmet; geoagmet.run(cfg)")
_table.add_row("cfg", "\\[geobase.txt, countries.txt, crops.txt, geocif.txt]")
_console.print(Panel(_table, title="[bold bright_white]GeoCIF Agmet Runner[/]", border_style="bright_blue", padding=(1, 2)))


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("DEFAULT", "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", [])
        eo_plot_raw = self._get_option(
            "eo_plot",
            ["ndvi", "cpc_tmax", "cpc_tmin", "chirps", "esi_4wk", "soil_moisture_as1", "soil_moisture_as2"],
        )
        self.eo_plot = self._expand_eo_plot(eo_plot_raw)
        self.eo_model = self._get_option(
            "eo_model",
            ["ndvi", "cpc_tmax", "cpc_tmin", "chirps", "esi_4wk", "soil_moisture_as1", "soil_moisture_as2"],
        )
        self.logo_harvest = self.dir_metadata / "images" / self.parser.get(
            "AGMET", "logo_harvest"
        )
        self.logo_geoglam = self.dir_metadata / "images" / self.parser.get(
            "AGMET", "logo_geoglam"
        )

    @staticmethod
    def _expand_eo_plot(eo_plot_raw):
        """Map raw EO variable names to subplot names in GEOGLAM display order.

        Layout (3 columns):
            Row 0: ndvi            | cumulative_precip | <soil/other>
            Row 1: yearly_ndvi     | daily_precip      | <soil/other>
            Row 2: esi_4wk         | cpc_tmax          | ...
        """
        # Desired order of subplot names matching GEOGLAM layout
        ordered = [
            "ndvi", "cumulative_precip", "nsidc_surface",
            "yearly_ndvi", "daily_precip", "nsidc_rootzone",
            "esi_4wk", "cpc_tmax",
        ]

        # Which raw vars expand into which subplot names
        expansions = {
            "ndvi": ["ndvi", "yearly_ndvi"],
            "chirps": ["cumulative_precip", "daily_precip"],
            "gcvi": ["gcvi", "yearly_gcvi"],
        }
        # These raw vars get dropped (absorbed into cpc_tmax avg temp plot)
        skip = {"cpc_tmin"}

        # Build set of subplot names from raw eo_plot
        subplot_set = set()
        for var in eo_plot_raw:
            if var in skip:
                continue
            if var in expansions:
                subplot_set.update(expansions[var])
            else:
                subplot_set.add(var)

        # Return in GEOGLAM order, then any extras not in the template
        result = [s for s in ordered if s in subplot_set]
        extras = [s for s in subplot_set if s not in ordered]
        return result + extras

    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_metadata / 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.add_yield_statistics(country, crop)

        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"
        )

    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)

        crop_short = utils.get_crop_abbrev(self.crop)

        # Abbreviate scale: admin_1 → adm1, admin_2 → adm2
        self.scale_short = self.scale.replace("admin_", "adm")

        folder = f"{crop_short}_s{self.growing_season}_{self.plot_season}"
        self.dir_agmet = (
            self.dir_output
            / "crop_condition"
            / ar.now().format("MMMM_DD_YYYY")
            / "plots"
            / 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:
            admin_level = self.parser.get(country, "admin_level")
            crops = ast.literal_eval(self.parser.get(country, "crops"))
            seasons = ast.literal_eval(
                self.parser.get(country, "seasons")
            )

            for crop in crops:
                for season in seasons:
                    all_combinations.append(
                        (country, admin_level, crop, 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 add_yield_statistics(self, country, crop):
        """Add yield data from FEWSNET/GEOGLAM stats to df_ccs."""
        # Add temporary columns expected by stats.add_statistics
        self.df_ccs["Region"] = self.df_ccs["region"]
        self.df_ccs["Harvest Year"] = self.df_ccs["harvest_season"]

        country_str = country.replace("_", " ").title()
        crop_str = utils.get_crop_name(crop)

        self.df_ccs = stats.add_statistics(
            dir_stats=self.dir_production_statistics,
            df=self.df_ccs,
            country=country_str,
            crop=crop_str,
            admin_zone=self.scale,
            stats=["Yield (tn per ha)"],
            method="",
            parser=self.parser,
        )

        # Map result back to the yield column used by agmet plotting
        if "Yield (tn per ha)" in self.df_ccs.columns:
            self.df_ccs["yield"] = self.df_ccs["Yield (tn per ha)"]

        # Drop temporary columns
        drop_cols = ["Region", "Harvest Year", "Yield (tn per ha)", "Area (ha)",
                     "Production (tn)", "Area"]
        self.df_ccs.drop(
            columns=[c for c in drop_cols if c in self.df_ccs.columns],
            inplace=True,
        )

    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):
        # Initialize chirps_gefs column
        self.df_ccs["chirps_gefs"] = np.nan
        self.df_region["chirps_gefs"] = 0.0

        try:
            bool_year_check, bool_date_check = self.check_date(
                self.df_region, plot_season
            )
        except Exception:
            return

        if not (bool_year_check & bool_date_check):
            return

        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"
        )

        region_name = self.df_region["region"].unique()[0]

        # Try current season, fall back to previous season
        gefs_files = list(path_gefs.glob(f"{region_name}*{plot_season}*.csv"))
        if not gefs_files:
            gefs_files = list(path_gefs.glob(f"{region_name}*{plot_season - 1}*.csv"))
        if not gefs_files:
            return

        df_gefs = pd.read_csv(gefs_files[0], header=None)

        # Remove leap year day (doy 60)
        if 60 in df_gefs.iloc[:, 4].values:
            df_gefs = df_gefs[df_gefs.iloc[:, 4] != 60]

        # Get all forecast values from column 5, replace NaN with 0
        val_gefs = np.nan_to_num(np.asarray(df_gefs.iloc[:, 5]))

        # Assign to 15-day forecast window (today → today+14)
        now = ar.utcnow().to("America/New_York")
        start_date = now.date()
        end_date = now.shift(days=+14).date()

        try:
            self.df_region.loc[start_date:end_date, "chirps_gefs"] = val_gefs
        except Exception:
            self.df_region.loc[start_date:end_date, "chirps_gefs"] = np.nan

        # Merge back into df_ccs via combine_first
        df_tmp = self.df_region.loc[start_date:end_date]
        self.df_region.set_index(
            self.df_region["datetime"] + self.df_region["region"],
            inplace=True, drop=True,
        )
        df_tmp.set_index(
            df_tmp["datetime"] + df_tmp["region"], inplace=True, drop=True
        )
        self.df_ccs.set_index(
            self.df_ccs["datetime"] + self.df_ccs["region"],
            inplace=True, drop=True,
        )
        self.df_ccs = self.df_ccs.combine_first(df_tmp)

        self.df_region = self.df_region.set_index(
            pd.to_datetime(self.df_region["datetime"]), drop=True
        )
        self.df_ccs = self.df_ccs.set_index(
            pd.to_datetime(self.df_ccs["datetime"]), drop=True
        )


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")
    df_crop_per_season.columns = df_crop_per_season.columns.str.lower()
    # Match crop: CSV uses abbreviations (mz, sb, ww) but obj.crop may be full name
    crop_key = utils.get_crop_abbrev(obj.crop)
    crop_name = df_crop_per_season[
        (df_crop_per_season["country"] == obj.country)
        & (df_crop_per_season["crop"] == crop_key)
        & (df_crop_per_season["season"] == int(obj.growing_season))
    ]["name"].values

    crop_name = crop_name[0] if len(crop_name) > 0 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)

    from geocif.data import ensure_metadata
    ensure_metadata(obj.parser)

    # 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()

    # Build and display run summary
    country_combos = {}
    for country, scale, crop, growing_season in all_combinations:
        country_combos.setdefault(country, []).append(
            f"{crop}, {scale}, s{growing_season}"
        )

    params = [("Countries", obj.countries)]
    for country, combos in country_combos.items():
        params.append((f"  {country}", "; ".join(dict.fromkeys(combos))))
    params.append(("Plot seasons", [str(s) for s in (obj.plot_seasons or [])]))
    params.append(("EO plot vars", obj.eo_plot))
    params.append(("EO model vars", obj.eo_model))
    params.append(("Total combinations", str(len(all_combinations))))
    ut.display_run_summary("GeoCIF Agmet Runner", params, wait=20)

    pbar = tqdm(all_combinations, total=len(all_combinations))
    for country, scale, crop, growing_season in pbar:
        pbar.set_description(f"{country} ({crop} s{growing_season})")
        # 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 tqdm(obj.plot_seasons, desc="Seasons", leave=False):
            # 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 tqdm(obj.list_regions, desc="Regions", leave=False):
                # 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 admin units
                ###############################################################
                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_short,
                    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",
                        ]

                try:
                    df_calendar_region = (
                        obj.df_region.groupby(
                            [
                                "country",
                                "calendar_region",
                                "harvest_season",
                                "doy",
                                "datetime",
                            ]
                        )[columns]
                        .mean()
                        .reset_index()
                    )
                except Exception:
                    continue
                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()
