"""Yield Outlook Map Generator.

Runs the ML pipeline for the current year and last N historical years,
then generates a diverging choropleth map showing current forecast yield
as a percentage of the historical mean/median prediction per region.
"""

import ast
import logging
import os
import sqlite3
import warnings
from pathlib import Path

import arrow as ar
import geopandas as gpd
import numpy as np
import palettable as pal
import pandas as pd

from geocif import geocif_runner as gc
from geocif import logger as log
from geocif import utils as ut
from .viz import plot

warnings.simplefilter(action="ignore", category=FutureWarning)

logger = logging.getLogger(__name__)

# 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 import yield_outlook; yield_outlook.run(cfg)")
_table.add_row("cfg", "\\[geobase.txt, countries.txt, crops.txt, geocif.txt]")
_console.print(
    Panel(
        _table,
        title="[bold bright_white]GeoCIF Yield Outlook[/]",
        border_style="bright_blue",
        padding=(1, 2),
    )
)


def _load_shapefiles(parser):
    """Load and concatenate shapefiles for all countries.

    Reuses the shapefile loading pattern from analysis.py Geoanalysis.setup().

    Returns:
        dg: Combined GeoDataFrame with 'Country Region' merge column.
        dict_config: Per country_crop config dict with method, crops, models, etc.
    """
    from geoprepare.georegion import get_boundary_col_mapping

    countries = ast.literal_eval(parser.get("DEFAULT", "countries"))
    dir_boundary_files = Path(parser.get("PATHS", "dir_boundary_files"))

    dict_config = {}
    all_shapefiles = []

    for country in countries:
        crops = ast.literal_eval(parser.get(country, "crops"))
        models = ast.literal_eval(parser.get(country, "models"))
        method = parser.get(country, "method")
        admin_zone = parser.get(country, "admin_level")

        for crop in crops:
            dict_config[f"{country}_{crop}"] = {
                "method": method,
                "crops": crop,
                "models": models,
                "admin_zone": admin_zone,
            }

        # Load shapefile
        shp_file = parser.get(country, "boundary_file")
        dg_country = gpd.read_file(
            dir_boundary_files / shp_file,
            engine="pyogrio",
        )

        # Rename columns using config-driven mapping
        rename = get_boundary_col_mapping(parser, shp_file)
        targets = set(rename.values())
        sources = set(rename.keys())
        conflicting = [
            c for c in dg_country.columns if c in targets and c not in sources
        ]
        if conflicting:
            dg_country = dg_country.drop(columns=conflicting)
        dg_country = dg_country.rename(columns=rename)

        if "ADM0_NAME" not in dg_country.columns:
            dg_country.loc[:, "ADM0_NAME"] = country.title().replace("_", " ")

        all_shapefiles.append(dg_country)

    dg = pd.concat(all_shapefiles, ignore_index=True)

    # Create "Country Region" merge column (same pattern as analysis.py:1229-1242)
    dg["Country Region"] = dg["ADM0_NAME"]
    dg["Country Region"] = dg["Country Region"].str.cat(dg["ADM1_NAME"], sep=" ")
    if "ADM2_NAME" in dg.columns:
        dg.loc[dg["ADM2_NAME"].notna(), "Country Region"] = (
            dg["ADM0_NAME"] + " " + dg["ADM2_NAME"]
        )
    dg["Country Region"] = dg["Country Region"].str.lower().replace("_", " ")

    return dg, dict_config


def _query_predictions(db_path, table, model, experiment_name="default"):
    """Query predictions from the SQLite database for a specific model.

    Returns DataFrame with columns: Country, Region, Harvest Year, Stage Name,
    Predicted Yield (tn per ha).
    """
    if not db_path.exists():
        logger.error("Database not found: %s", db_path)
        return pd.DataFrame()

    con = sqlite3.connect(db_path)
    try:
        df = pd.read_sql(
            f'SELECT "Country", "Region", "Harvest Year", "Stage Name", '
            f'"Predicted Yield (tn per ha)" '
            f'FROM "{table}" WHERE "Experiment Name" = ? AND "Model" = ?',
            con,
            params=(experiment_name, model),
        )
    except (pd.errors.DatabaseError, sqlite3.OperationalError) as e:
        logger.warning("Failed to query table '%s': %s", table, e)
        df = pd.DataFrame()
    finally:
        con.close()
    return df


def _compute_outlook_index(df, current_year, n_years, aggregation, stage_name):
    """Compute yield outlook index per region.

    outlook_index = (current_predicted / agg(historical_predicted)) * 100

    Returns DataFrame with columns: Country, Region, Country Region,
    current_predicted, hist_predicted, outlook_index.
    """
    df_stage = df[df["Stage Name"] == stage_name].copy()

    # Current year predictions per region
    df_current = df_stage[df_stage["Harvest Year"] == current_year]
    current_pred = (
        df_current.groupby(["Country", "Region"])["Predicted Yield (tn per ha)"]
        .mean()
        .rename("current_predicted")
    )

    # Historical years per region
    min_year = current_year - n_years
    df_hist = df_stage[
        (df_stage["Harvest Year"] < current_year)
        & (df_stage["Harvest Year"] >= min_year)
    ]
    agg_func = "median" if aggregation == "median" else "mean"
    hist_agg = (
        df_hist.groupby(["Country", "Region"])["Predicted Yield (tn per ha)"]
        .agg(agg_func)
        .rename("hist_predicted")
    )

    # Compute index
    df_outlook = pd.concat([current_pred, hist_agg], axis=1).dropna()
    df_outlook["outlook_index"] = np.where(
        df_outlook["hist_predicted"] != 0,
        df_outlook["current_predicted"] / df_outlook["hist_predicted"] * 100,
        np.nan,
    )
    df_outlook = df_outlook.reset_index()

    # Create merge column (same pattern as analysis.py:1410-1414)
    df_outlook["Country Region"] = (
        df_outlook["Country"].str.lower().str.replace("_", " ")
        + " "
        + df_outlook["Region"].str.lower()
    )

    return df_outlook


def _generate_outlook_map(
    dg,
    df_outlook,
    countries,
    crop,
    model,
    current_year,
    n_years,
    aggregation,
    dir_out,
    stage_name="",
    annotate_regions=False,
):
    """Generate a diverging choropleth map of the yield outlook index."""
    col = "outlook_index"

    # Symmetric range around 100
    deviation = max(
        abs(df_outlook[col].max() - 100),
        abs(100 - df_outlook[col].min()),
    )
    deviation = max(deviation, 10)  # at least +/-10
    vmin = 100 - deviation
    vmax = 100 + deviation

    countries_display = [c.title().replace("_", " ") for c in countries]
    stage_suffix = f"_{stage_name}" if stage_name else ""
    fname = f"yield_outlook_{'_'.join(countries)}_{crop}_{model}{stage_suffix}_{current_year}.png"

    stage_label = f", stage {stage_name}" if stage_name else ""
    plot.plot_df_shpfile(
        dg,
        df_outlook,
        merge_col="Country Region",
        name_country=countries_display,
        name_col=col,
        dir_out=dir_out,
        fname=fname,
        label=f"% of {n_years}-yr {aggregation}\npredicted yield ({crop.title()}, {current_year}{stage_label})",
        vmin=vmin,
        vmax=vmax,
        cmap=pal.colorbrewer.diverging.BrBG_11,
        series="diverging",
        show_bg=False,
        annotate_regions=annotate_regions,
        loc_legend="lower left",
    )


def run(path_config_files=None, current_year=None, n_years=None, aggregation=None):
    """Main entry point for yield outlook map generation.

    1. Override forecast_seasons to cover [current_year - n_years, ..., current_year]
    2. Run the ML pipeline via gc.execute_models()
    3. Query predictions from the database
    4. Compute outlook index per region
    5. Generate diverging choropleth maps and CSV

    Args:
        path_config_files: List of config file paths.
        current_year: Forecast year (default: this year).
        n_years: Number of historical years for comparison (default: config or 10).
        aggregation: 'mean' or 'median' (default: config or 'mean').
    """
    if path_config_files is None:
        path_config_files = [Path("../config/geocif.txt")]

    logger_obj, parser = log.setup_logger_parser(path_config_files)

    # Read config with defaults
    if n_years is None:
        n_years = parser.getint("ML", "outlook_n_years", fallback=10)
    if aggregation is None:
        aggregation = parser.get("ML", "outlook_aggregation", fallback="mean")
    if current_year is None:
        current_year = ar.utcnow().to("America/New_York").year

    countries = ast.literal_eval(parser.get("DEFAULT", "countries"))
    experiment_name = parser.get("DEFAULT", "experiment_name", fallback="default")

    # ---- Step 1: Run ML pipeline for all needed years ----
    outlook_seasons = list(range(current_year - n_years, current_year + 1))
    originals = {}
    for country in countries:
        originals[country] = parser.get(country, "forecast_seasons")
        parser.set(country, "forecast_seasons", str(outlook_seasons))

    parser.set("DEFAULT", "experiment_name", "outlook")
    orig_parallel = parser.get("DEFAULT", "do_parallel_ml", fallback=None)
    parser.set("DEFAULT", "do_parallel_ml", "True")

    inputs = gc.gather_inputs(parser)

    params = [
        ("Countries", countries),
        ("Forecast year", str(current_year)),
        ("Lookback years", str(n_years)),
        ("Seasons", f"{outlook_seasons[0]}-{outlook_seasons[-1]}"),
        ("Aggregation", aggregation),
        ("DB", parser.get("DEFAULT", "db")),
        ("Total combinations", str(len(inputs))),
    ]
    ut.display_run_summary("GeoCIF Yield Outlook", params, wait=10)

    gc.execute_models(inputs, logger_obj, parser)

    # Restore original config values
    for country, orig in originals.items():
        parser.set(country, "forecast_seasons", orig)
    parser.set("DEFAULT", "experiment_name", experiment_name)
    if orig_parallel is not None:
        parser.set("DEFAULT", "do_parallel_ml", orig_parallel)
    elif parser.has_option("DEFAULT", "do_parallel_ml"):
        parser.remove_option("DEFAULT", "do_parallel_ml")

    # ---- Step 2: Load shapefiles ----
    dg, dict_config = _load_shapefiles(parser)

    # ---- Step 3: Query DB, compute outlook, generate maps ----
    project_name = parser.get("DEFAULT", "project_name", fallback="geocif")
    dir_output = Path(parser.get("PATHS", "dir_output")) / project_name
    db_name = parser.get("DEFAULT", "db")
    db_path = dir_output / "ml" / "db" / db_name

    today = ar.utcnow().to("America/New_York").format("MMMM_DD_YYYY")
    dir_outlook = dir_output / "ml" / "analysis" / today / "outlook"
    os.makedirs(dir_outlook, exist_ok=True)

    all_outlook_frames = []

    for country_crop, config in dict_config.items():
        crop = config["crops"]
        country = country_crop.replace(f"_{crop}", "")
        models = config["models"]

        for model in models:
            logger.info("Yield outlook: %s %s %s", country, crop, model)

            df = _query_predictions(db_path, country_crop, model, experiment_name="outlook")
            if df.empty:
                logger.warning("No predictions found for %s %s %s", country, crop, model)
                continue

            # Get all stages available for current year
            df_current = df[df["Harvest Year"] == current_year]
            if df_current.empty:
                logger.warning(
                    "No predictions for year %d in %s %s %s",
                    current_year, country, crop, model,
                )
                continue

            annotate = parser.getboolean(
                country, "annotate_regions", fallback=False
            )

            for stage_name in df_current["Stage Name"].unique():
                # Compute outlook index
                df_outlook = _compute_outlook_index(
                    df, current_year, n_years, aggregation, stage_name
                )
                if df_outlook.empty:
                    logger.warning(
                        "Could not compute outlook for %s %s %s stage %s",
                        country, crop, model, stage_name,
                    )
                    continue

                n_hist = len(
                    df[
                        (df["Stage Name"] == stage_name)
                        & (df["Harvest Year"] < current_year)
                        & (df["Harvest Year"] >= current_year - n_years)
                    ]["Harvest Year"].unique()
                )
                if n_hist < 3:
                    logger.warning(
                        "Only %d historical years for %s %s %s stage %s (requested %d)",
                        n_hist, country, crop, model, stage_name, n_years,
                    )

                df_outlook["Crop"] = crop
                df_outlook["Model"] = model
                df_outlook["Stage Name"] = stage_name
                df_outlook["Forecast Year"] = current_year
                all_outlook_frames.append(df_outlook)

                # Generate map
                _generate_outlook_map(
                    dg,
                    df_outlook,
                    [country],
                    crop,
                    model,
                    current_year,
                    n_years,
                    aggregation,
                    dir_outlook,
                    stage_name=stage_name,
                    annotate_regions=annotate,
                )
                logger.info(
                    "Map saved: %s",
                    dir_outlook
                    / f"yield_outlook_{country}_{crop}_{model}_{stage_name}_{current_year}.png",
                )

    # Save combined CSV
    if all_outlook_frames:
        df_all = pd.concat(all_outlook_frames, ignore_index=True)
        csv_path = dir_outlook / f"yield_outlook_{current_year}.csv"
        df_all.to_csv(csv_path, index=False)
        logger.info("Outlook CSV saved to %s", csv_path)
    else:
        logger.warning("No outlook data generated — check DB has predictions.")


if __name__ == "__main__":
    run()
