"""
IMF World Economic Outlook data fetching and processing.

This module provides the IMFData class for fetching and processing
data from the IMF World Economic Outlook database.

Example usage:
    from atlas_common_data.sources.imf import IMFData

    imf = IMFData()

    # Get specific indicators
    df = imf.get_indicators(['NGDPD', 'LP'], country_codes=['USA', 'CHN'])

    # Get all standard indicators
    df = imf.get_all_indicators()

    # Get GDP data specifically
    df = imf.get_gdp_data()
"""

from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Union
import warnings

import json
import pandas as pd
import requests


class IMFData:
    """Fetch and process IMF World Economic Outlook data.

    This class provides methods to query the IMF API and process the returned
    data into analysis-ready DataFrames. Unlike the legacy implementation,
    this class returns DataFrames without writing files to disk.

    Attributes:
        IMF_INDICATORS: Mapping of IMF indicator codes to friendly names
        MULTIPLIERS: Unit conversion multipliers for certain indicators
        INCOME_LEVEL_PERCENTILES: Percentile thresholds for income classification

    Example:
        >>> imf = IMFData()
        >>> df = imf.get_indicators(['NGDPD'])
        >>> print(df.columns.tolist())
        ['year', 'iso3_code', 'gdp']
    """

    # Mapping of IMF indicator codes to friendly column names
    IMF_INDICATORS: Dict[str, str] = {
        "NGDP_R": "gdp_const",
        "NGDPD": "gdp",
        "NGDP_RPCH": "gdp_const_growth",
        "PPPGDP": "gdp_ppp",
        "NGDPDPC": "gdppc",
        "PPPPC": "gdppc_ppp",
        "LP": "population",
        "BCA": "current_account",
        "PCPIPCH": "inflation_rate",
    }

    # Multipliers for unit conversion (some IMF values are in millions/billions)
    MULTIPLIERS: Dict[str, int] = {
        "NGDPD": 1_000_000_000,
        "PPPGDP": 1_000_000_000,
        "LP": 1_000_000,
        "BCA": 1_000_000_000,
    }

    # Percentile thresholds for income level classification
    INCOME_LEVEL_PERCENTILES: Dict[str, float] = {
        "low": 0,
        "lower middle": 0.12,
        "upper middle": 0.39,
        "high": 0.66,
    }

    # Default indicators for comprehensive data fetch
    DEFAULT_INDICATORS: List[str] = [
        "NGDP_R",  # GDP, constant prices
        "NGDP_RPCH",  # GDP growth rate
        "NGDPD",  # GDP, current prices (USD)
        "PPPGDP",  # GDP, PPP
        "NGDPDPC",  # GDP per capita, current prices
        "PPPPC",  # GDP per capita, PPP
        "LP",  # Population
        "BCA",  # Current account balance
    ]

    # IMF API base URL
    API_BASE_URL = "https://www.imf.org/external/datamapper/api/v1"

    def __init__(
        self,
        latest_year: Optional[int] = None,
    ):
        """Initialize IMFData.

        Args:
            latest_year: Default end year for all queries. Methods will use
                this when end_year/reference_year is not explicitly passed.
        """
        self.latest_year = latest_year
        self._classification_dir = (
                Path(__file__).parent.parent / "data" / "classification"
            )

    def _query_api(
        self,
        indicators: List[str],
        country_codes: Optional[List[str]] = None,
    ) -> pd.DataFrame:
        """Query the IMF API for specified indicators.

        Args:
            indicators: List of IMF indicator codes (e.g., ['NGDPD', 'LP'])
            country_codes: Optional list of ISO3 country codes to filter

        Returns:
            DataFrame with columns: year, iso3_code, and one column per indicator
            (using friendly names from IMF_INDICATORS mapping)

        Raises:
            requests.RequestException: If API request fails
            ValueError: If response cannot be parsed
        """
        if not indicators:
            raise ValueError(
                "indicators list cannot be empty. "
                "Pass at least one indicator code, e.g. ['NGDPD']."
            )

        fields_str = "/".join(indicators)
        country_str = "/".join(country_codes) if country_codes else ""

        url = f"{self.API_BASE_URL}/{fields_str}/{country_str}"
        try:
            response = requests.get(url, timeout=60)
            response.raise_for_status()
            data = response.json()
        except requests.exceptions.RequestException as e:
            raise RuntimeError(
                f"IMF API request failed for indicators {indicators}: {e}"
            ) from e
        except json.JSONDecodeError as e:
            raise RuntimeError(
                f"IMF API returned invalid JSON for indicators {indicators}: {e}"
            ) from e

        if "values" not in data or not data["values"]:
            return pd.DataFrame()

        # Parse the nested JSON response
        df = pd.json_normalize(data["values"]).T.reset_index()

        # Split the index into components
        df[["indicator", "iso3_code", "year"]] = df["index"].str.split(".", expand=True)
        df = (
            df.drop(columns="index")
            .rename(columns={0: "value"})
            .astype({"indicator": str, "iso3_code": str, "year": int, "value": float})
        )

        # Pivot to wide format
        df = df.pivot_table(
            values="value", index=["year", "iso3_code"], columns="indicator"
        ).reset_index()

        # Apply multipliers for unit conversion
        for col in df.columns:
            if col in self.MULTIPLIERS:
                df[col] = df[col] * self.MULTIPLIERS[col]

        # Rename columns to friendly names
        df = df.rename(columns=self.IMF_INDICATORS, errors="ignore")

        return df

    def _load_country_classification(self) -> pd.DataFrame:
        """Load country classification data.

        Returns:
            DataFrame with country_id and iso3_code columns
        """
        path = self._classification_dir / "location" / "country.csv"
        return pd.read_csv(path)[["country_id", "iso3_code"]]

    def _load_in_country_profiles(self) -> set:
        """Load the set of iso3 codes for countries that are in country profiles.

        Returns:
            Set of iso3_code strings where in_cp is True.
        """
        path = self._classification_dir.parent / "classification" / "location" / "country.csv"
        df = pd.read_csv(path)
        return set(df.loc[df["in_cp"] == True, "iso3_code"])

    def _load_region_classification(self) -> pd.DataFrame:
        """Load region classification data.

        Returns:
            DataFrame with country_id and region columns
        """
        path = self._classification_dir / "location" / "group_member.csv"
        df = pd.read_csv(path)
        df = df[df.group_type == "region"].rename(columns={"group_name": "region"})
        return df[["country_id", "region"]]

    def _merge_country_id(self, df: pd.DataFrame) -> pd.DataFrame:
        """Merge Atlas country_id into DataFrame.

        Args:
            df: DataFrame with iso3_code column

        Returns:
            DataFrame with country_id column added
        """
        countries = self._load_country_classification()
        return df.merge(countries, on="iso3_code", how="inner")

    def get_indicators(
        self,
        indicators: Optional[List[str]] = None,
        country_codes: Optional[List[str]] = None,
        start_year: Optional[int] = None,
        end_year: Optional[int] = None,
        include_country_id: bool = True,
    ) -> pd.DataFrame:
        """Fetch IMF indicators and return as DataFrame.

        This is the primary method for fetching IMF data. It queries the API,
        applies unit conversions, and optionally filters by year range.

        Args:
            indicators: List of IMF indicator codes. If None, uses DEFAULT_INDICATORS.
                Available indicators:
                - NGDP_R: GDP, constant prices
                - NGDP_RPCH: GDP growth rate
                - NGDPD: GDP, current prices (USD)
                - PPPGDP: GDP, PPP
                - NGDPDPC: GDP per capita, current prices
                - PPPPC: GDP per capita, PPP
                - LP: Population
                - BCA: Current account balance
                - PCPIPCH: Inflation rate
            country_codes: Optional list of ISO3 country codes to filter
            start_year: Optional minimum year to include
            end_year: Optional maximum year to include
            include_country_id: If True, merge Atlas country_id from classification

        Returns:
            DataFrame with columns:
                - year: Year of observation
                - iso3_code: ISO 3166-1 alpha-3 country code
                - country_id: Atlas country ID (if include_country_id=True)
                - Plus one column per indicator (using friendly names)

        Example:
            >>> imf = IMFData()
            >>> df = imf.get_indicators(
            ...     indicators=['NGDPD', 'LP'],
            ...     country_codes=['USA', 'CHN'],
            ...     start_year=2015,
            ...     end_year=2023
            ... )
        """
        if indicators is None:
            indicators = self.DEFAULT_INDICATORS.copy()

        # Apply latest_year default
        if end_year is None and self.latest_year is not None:
            end_year = self.latest_year

        df = self._query_api(indicators, country_codes)

        if df.empty:
            return df

        # Filter by year range if specified
        if start_year is not None:
            df = df[df["year"] >= start_year]
        if end_year is not None:
            df = df[df["year"] <= end_year]

        # Sort by country and year
        df = df.sort_values(["iso3_code", "year"])

        # Optionally merge country_id
        if include_country_id:
            df = self._merge_country_id(df)

        return df

    def get_inflation_index(
        self,
        base_year: int,
        country_code: str = "USA",
    ) -> pd.DataFrame:
        """Calculate inflation index from IMF inflation rate data.

        Args:
            base_year: Year to use as base (index = 1.0)
            country_code: Country to use for inflation (default: USA)

        Returns:
            DataFrame with columns:
                - year
                - inflation_rate
                - inflation_index (normalized to base_year)
        """
        df = self._query_api(["PCPIPCH"], [country_code])

        if df.empty:
            return df

        df = df.drop(columns="iso3_code", errors="ignore")

        # Initialize and calculate cumulative index
        df = df.sort_values("year").reset_index(drop=True)
        df["inflation_index"] = 100.0

        for i in range(1, len(df)):
            df.loc[i, "inflation_index"] = df.loc[i - 1, "inflation_index"] * (
                1 + df.loc[i, "inflation_rate"] / 100
            )

        # Normalize to base year
        base_value = df.loc[df["year"] == base_year, "inflation_index"]
        if len(base_value) > 0:
            df["inflation_index"] = df["inflation_index"] / base_value.values[0]

        return df

    def get_average_growth_rates(
        self,
        lookback_periods: List[int] = [3, 5, 10, 15],
        reference_year: Optional[int] = None,
        country_codes: Optional[List[str]] = None,
    ) -> pd.DataFrame:
        """Calculate average growth rates over lookback periods.

        Calculates GDP per capita growth rates averaged over specified
        lookback periods (e.g., 3-year, 5-year, 10-year averages).

        Args:
            lookback_periods: List of years to look back for averaging
            reference_year: Year to report averages for. If None, uses current year - 1
            country_codes: Optional list of ISO3 country codes to filter

        Returns:
            DataFrame with columns:
                - year: Reference year
                - iso3_code: Country code
                - country_id: Atlas country ID
                - lookback: Lookback period
                - gdppc_const_growth: Average growth rate
        """
        if reference_year is None:
            reference_year = (
                self.latest_year
                if self.latest_year is not None
                else datetime.now().year - 1
            )

        df = self._query_api(["NGDP_RPCH", "LP"], country_codes)

        if df.empty:
            return df

        df = df.sort_values(["iso3_code", "year"])

        # Convert growth rate to decimal
        df["gdp_const_growth"] = df["gdp_const_growth"] / 100.0

        # Calculate population growth
        df["pop_growth"] = df.groupby("iso3_code")["population"].ffill().pct_change(fill_method=None)

        # Calculate GDP per capita growth
        df["gdppc_const_growth"] = df["gdp_const_growth"] - df["pop_growth"]

        # Merge country_id
        df = self._merge_country_id(df)

        # Calculate rolling averages for each lookback period
        for period in lookback_periods:
            df[f"gdppc_const_growth_{period}"] = df.groupby("iso3_code")[
                "gdppc_const_growth"
            ].transform(lambda x: x.rolling(period).mean())

        # Filter to reference year
        df = df[df["year"] == reference_year].drop(
            columns=[
                "gdp_const_growth",
                "population",
                "pop_growth",
                "gdppc_const_growth",
            ],
            errors="ignore",
        )

        # Reshape to long format
        df = pd.wide_to_long(
            df,
            stubnames="gdppc_const_growth_",
            i=["year", "iso3_code", "country_id"],
            j="lookback",
        ).reset_index()

        df = df.rename(columns={"gdppc_const_growth_": "gdppc_const_growth"})
        df = df.astype({"year": int, "lookback": int})

        return df

    def get_country_metadata(
        self,
        reference_year: Optional[int] = None,
    ) -> pd.DataFrame:
        """Get country-level metadata including income levels.

        Fetches data for a single year and classifies countries by income level
        based on GDP per capita percentiles.

        Args:
            reference_year: Year to use for classification. If None, uses current year - 1

        Returns:
            DataFrame with columns:
                - country_id: Atlas country ID
                - iso3_code: Country code
                - region: Geographic region
                - incomelevel_enum: Income level classification
        """
        if reference_year is None:
            reference_year = (
                self.latest_year
                if self.latest_year is not None
                else datetime.now().year - 1
            )

        # Get GDP per capita data
        df = self.get_indicators(
            indicators=["NGDPDPC"],
            start_year=reference_year,
            end_year=reference_year,
        )

        if df.empty:
            return df

        # Classify by income level using percentiles
        df["incomelevel_enum"] = pd.qcut(
            df["gdppc"],
            q=list(self.INCOME_LEVEL_PERCENTILES.values()) + [1],
            labels=list(self.INCOME_LEVEL_PERCENTILES.keys()),
        )

        # Merge region data
        regions = self._load_region_classification()
        df = df.merge(regions, on="country_id", how="left")

        return df[["country_id", "iso3_code", "region", "incomelevel_enum"]]

    def get_imf_data(
        self,
        end_year: int,
    ) -> pd.DataFrame:
        """Fetch IMF indicators and apply FRED PPIIDC deflator to produce constant-dollar columns.

        This is the primary data method matching the output of the legacy
        ``calculate_imf_data()`` function. It produces the same columns as
        ``imf_indicators/data/imf_data.csv``.

        Constant-dollar columns (suffixed ``_const``) are computed by dividing each
        nominal value by the FRED Producer Price Index by Commodity: Industrial
        Commodities (PPIIDC) deflator, fetched from the FRED API via ``FREDData``.
        The deflator is normalised to 1.0 at ``end_year``, so all constant values
        are expressed in ``end_year`` dollars.

        Columns produced by the deflation step:
            gdp_const, gdp_ppp_const, gdppc_const, gdppc_ppp_const

        Args:
            end_year: Latest year to include. Also used as the FRED PPIIDC deflator
                base year (i.e. constant-dollar values are in ``end_year`` dollars).
                Defaults to ``self.latest_year`` if set.

        Returns:
            DataFrame with columns:
                year, iso3_code, country_id,
                current_account, population,
                gdp, gdppc, gdp_const_growth,
                gdp_ppp, gdppc_ppp, gdp_const,
                gdp_ppp_const, gdppc_const, gdppc_ppp_const
        """
        from atlas_common_data.sources.fred import FREDData

        if end_year is None:
            end_year = self.latest_year

        df = self.get_indicators(end_year=end_year)

        if df.empty:
            return df

        # Warn if any countries are missing data for the requested end_year
        if end_year is not None:
            all_countries = set(df["iso3_code"].unique())
            countries_in_latest = set(df[df["year"] == end_year]["iso3_code"].unique())
            missing = sorted(all_countries - countries_in_latest)
            if missing:
                in_country_profiles = self._load_in_country_profiles()
                missing_ranked = sorted(c for c in missing if c in in_country_profiles)
                if missing_ranked:
                    warnings.warn(
                        f"{len(missing_ranked)} ranked countries have no IMF data for {end_year} "
                        f"and may not have reported yet: {missing_ranked}",
                        UserWarning,
                        stacklevel=2,
                    )

        # Columns that get constant-dollar equivalents (mirrors old code logic)
        _skip = {"gdp_const", "gdp_const_growth", "population", "current_account"}
        nominal_cols = [
            c
            for c in df.columns
            if c in set(self.IMF_INDICATORS.values()) and c not in _skip
        ]

        # Apply FRED deflator
        fred = FREDData()
        deflator = fred.get_inflation_index(base_year=end_year)[["year", "deflator"]]
        df["year"] = df["year"].astype(int)
        df = df.merge(deflator, on="year", how="left")
        for col in nominal_cols:
            df[f"{col}_const"] = df[col] / df["deflator"]
        df = df.drop(columns=["deflator"])

        return df.sort_values(["iso3_code", "year"])

    def get_gdp_data(
        self,
        start_year: Optional[int] = None,
        end_year: Optional[int] = None,
    ) -> pd.DataFrame:
        """Fetch GDP-related indicators.

        Convenience method that fetches all GDP, GDP per capita, and PPP indicators.

        Args:
            start_year: Optional minimum year to include.
            end_year: Optional maximum year to include.

        Returns:
            DataFrame with GDP-related columns.
        """
        gdp_indicators = ["NGDPD", "NGDPDPC", "NGDP_R", "NGDP_RPCH", "PPPGDP", "PPPPC"]
        return self.get_indicators(
            indicators=gdp_indicators,
            start_year=start_year,
            end_year=end_year,
        )

    @classmethod
    def list_available_indicators(cls) -> Dict[str, str]:
        """List available IMF indicator codes and their descriptions.

        Returns:
            Dictionary mapping IMF codes to friendly names
        """
        return cls.IMF_INDICATORS.copy()
