import numpy as np
import pandas as pd
from numpy import ndarray
from pandas import DataFrame

from sarawater.scenarios import Scenario, ConstScenario
from sarawater.IHA import compute_IHA
from sarawater.utils import _validate_positive_numeric


class Reach:
    def __init__(self, name: str, dates: list, Qnat: ndarray, Qabs_max: float):
        """Represents a river reach.

        Parameters
        ----------
        name : str
            Name of the reach.
        dates : list[datetime]
            List of dates for the time series.
        Qnat : ndarray
            Natural flow rate time series.
        Qabs_max : float
            Maximum value for the water abstraction.
        """
        if len(dates) != len(Qnat):
            raise ValueError("Dates and flow data length mismatch")

        # Validate that all discharge values are non-negative
        if np.any(Qnat < 0):
            raise ValueError("Natural flow rate (Qnat) must be non-negative (>= 0)")

        # Validate Qabs_max is positive
        _validate_positive_numeric(Qabs_max, "Qabs_max")

        self.name = name
        self.dates = dates
        self.Qnat = Qnat
        self.Qabs_max = Qabs_max
        self.scenarios: list[Scenario] = []
        self.IHA_nat = compute_IHA(Qnat, Qnat, dates)

    def __str__(self):
        return f"{self.name} is a Reach object with a flow time series with {len(self.Qnat)} elements. The date range starts from {min(self.dates)} and has {len(self.dates)} elements. The maximum flow abstraction is Qabs_max={self.Qabs_max} m3/s. So far, {len(self.scenarios)} scenarios have been added."

    def add_scenario(self, scenario: Scenario):
        """Add a scenario to the reach.

        Parameters
        ----------
        scenario : Scenario
            The scenario to add.

        Returns
        -------
        Reach
            The current reach instance.
        """
        self.scenarios.append(scenario)
        return self

    def print_scenarios(self):
        """Print the list of scenarios added to the reach."""
        for i, scenario in enumerate(self.scenarios):
            print(f"scenarios[{i}]: {scenario.name} | {scenario.description}")
        return None

    def add_ecological_flow_scenario(
        self, name: str, description: str, k: float = 0.2, p: float = 2.0
    ) -> Scenario:
        """Add an ecological flow scenario to the reach.

        Parameters
        ----------
        name : str
            Name of the scenario
        description : str
            Description of the scenario
        k : float, optional
            Protection factor, by default 0.2
        p : float, optional
            Nature protection factor, by default 2.0

        Returns
        -------
        Scenario
            The created ecological flow scenario
        """
        # Calculate monthly averages from daily data
        monthly_means = np.zeros(12)
        for month in range(1, 13):
            # Create mask for current month across all years
            month_mask = np.array([d.month == month for d in self.dates])
            if np.any(month_mask):  # Check if we have data for this month
                monthly_means[month - 1] = np.mean(self.Qnat[month_mask])
            else:
                raise ValueError(f"No data available for month {month}")

        # Calculate overall mean flow from daily data
        Q_mean = np.mean(self.Qnat)

        # Calculate Q97 from daily data (approximation of Q355)
        Q97 = np.percentile(self.Qnat, 3)

        # Calculate DE for each month
        Qreq_months = []
        for Q_month in monthly_means:
            M1 = np.sqrt(Q_month / Q_mean)
            DE = k * p * M1 * Q_mean
            Qreq_months.append(max(DE, Q97))

        # Create and return constant scenario with computed monthly values
        scenario = ConstScenario(name, description, self, Qreq_months)
        self.add_scenario(scenario)
        return scenario

    def add_HQ_curve(self, HQ_curve: DataFrame):
        """Add a habitat-flow curve to the reach.

        Parameters
        ----------
        HQ_curve : DataFrame
            Habitat-flow curve as a pandas DataFrame where columns represent different species/stages.

        Returns
        -------
        Reach
            The current reach instance.
        """
        self.HQ_curve = HQ_curve
        self.HQ_curve_columns = list(HQ_curve.columns)
        # Exclude "DIS" and "WET" columns from available curves
        self.available_HQ_curves = [
            col for col in self.HQ_curve_columns if col not in ["DIS", "WET"]
        ]
        return self

    def get_list_available_HQ_curves(self) -> list:
        """Get the list of available HQ curve names.

        Returns
        -------
        list
            List of column names representing available HQ curves.
        """
        if hasattr(self, "available_HQ_curves"):
            return self.available_HQ_curves
        else:
            return []

    def get_HQ_curve(self, curve_name: str) -> DataFrame:
        """Get a specific habitat-flow curve by name.

        Parameters
        ----------
        curve_name : str
            Name of the desired HQ curve (species).

        Returns
        -------
        DataFrame
            The requested HQ curve as a pandas DataFrame.
        """
        if hasattr(self, "HQ_curve"):
            if curve_name in self.HQ_curve_columns:
                HQ_Q = self.HQ_curve["DIS"]
                HQ_H = self.HQ_curve[curve_name]
                HQ = pd.DataFrame({"DIS": HQ_Q, curve_name: HQ_H})
                return HQ
            else:
                raise ValueError(f"HQ curve '{curve_name}' not found in the reach.")
        else:
            raise ValueError("No HQ curves have been added to this reach.")

    def add_cross_section_geometry(
        self,
        slope,
        ks,
        width=None,
        section=None,
    ):
        """Add cross-section geometry, channel roughness and bed slope to the reach. Either coordinate pairs (composite cross section) or channel width (rectangular cross section) must be provided. When a width is provided, a two-point simple rectangular cross-section is created with a flat bed at elevation 0, corresponding to the coordinate pairs y [m] = [0, width] and z [m] = [0, 0]. When section coordinates are provided, they must be in the form of a CSV file or a DataFrame with columns 'y [m]' (transverse coordinate) and 'z [m]' (bed elevation).

        Parameters
        ----------
        slope : float
            Channel bed slope (m/m).

        ks : float
            Strickler coefficient (m^(1/3)/s) representing channel roughness.

        width : float, optional
            Channel width for rectangular cross-section (meters).

        section : str or DataFrame
            Path to CSV file containing section coordinates ('y [m]', 'z [m]'),
            or a pandas DataFrame with those columns. y represents transverse coordinates along the cross-section, and z represents bed elevation at each point.

        Returns
        -------
        Reach
            The current reach instance.
        """
        if width is None and section is None:
            raise ValueError(
                "Either width for rectangular section or section coordinates must be provided."
            )
        if width is not None and section is not None:
            raise ValueError(
                "Provide either width for rectangular section or section coordinates, not both."
            )

        _validate_positive_numeric(slope, "slope")
        _validate_positive_numeric(ks, "ks")

        self.ks = ks
        self.slope = slope

        # Validate and process width or section input to retrieve the (y,z) coordinates of the cross-section
        if width is not None:
            _validate_positive_numeric(width, "width")

            # Create simple rectangular section data
            y = np.array([0, width])
            z = np.array([0, 0])  # Flat bed at elevation 0
            cross_section_coordinates = pd.DataFrame({"y [m]": y, "z [m]": z})

        if section is not None:
            if isinstance(section, str):
                cross_section_coordinates = pd.read_csv(section, delimiter=None)
            elif isinstance(section, pd.DataFrame):
                cross_section_coordinates = section.copy()
            else:
                raise ValueError(
                    "section must be a CSV file path or a pandas DataFrame"
                )

            # Validate columns
            required_cols = ["y [m]", "z [m]"]
            if not all(
                col in cross_section_coordinates.columns for col in required_cols
            ):
                raise ValueError(f"Section data must contain columns {required_cols}")

            # Validate numeric values in section
            if not np.all(np.isfinite(cross_section_coordinates["y [m]"])):
                raise ValueError("Section y coordinates must be finite numbers")
            if not np.all(np.isfinite(cross_section_coordinates["z [m]"])):
                raise ValueError(
                    "Section z coordinates (representing bed elevation) must be finite numbers"
                )

            if not np.all(np.diff(cross_section_coordinates["y [m]"]) > 0):
                raise ValueError("y coordinates must be monotonically increasing")

        self.cross_section_coordinates = cross_section_coordinates
        return self

    def add_grain_size_distribution(self, grain_data):
        """Add grain size distribution data to the reach. The input can be a single D50 value, a DataFrame with columns 'i(di)' and 'di[mm]', a 2D array with those columns, or a path to a CSV file containing that data.

        Parameters
        ----------
        grain_data : float or DataFrame or array or str
            Grain size distribution data. Can be one of the following:
            - A single float representing D50 in millimeters. This will be converted to a uniform grain size distribution with 100% in the corresponding phi class.
            - A pandas DataFrame with columns 'i(di)' (cumulative percentage) and 'di[mm]' (diameter in millimeters).
            - A 2D array with columns 'i(di)' and 'di[mm]'.
            - A path to a CSV file containing columns 'i(di)' and 'di[mm]'.

        Returns
        -------
        Reach
            The current reach instance for method chaining.
        """
        # Define standard phi classes: -9.5 to 7.5 with step 1 (18 classes total)
        phi_class_centers = np.arange(-9.5, 7.5 + 1, 1)
        expected_phi_length = len(phi_class_centers)  # Should be 18

        # Initialize variables that will be set in all branches
        dfphi = None
        phi_percentages = None

        if isinstance(grain_data, (float, int)):
            # Single D50 value provided - assign 100% to the corresponding phi class
            d50 = float(grain_data)
            if d50 <= 0:
                raise ValueError("D50 must be positive")

            # Convert D50 to phi scale
            phi_d50 = -np.log2(d50)

            # Find the closest phi class center
            # phi_class_centers = [-9.5, -8.5, -7.5, ..., 6.5, 7.5]
            closest_phi_idx = np.argmin(np.abs(phi_class_centers - phi_d50))
            closest_phi = phi_class_centers[closest_phi_idx]

            # Create phi_percentages with 100% in the closest class, 0% elsewhere
            phi_percentages = pd.Series(0.0, index=phi_class_centers)
            phi_percentages.loc[closest_phi] = 1.0

            # dfphi remains None for D50 input

        else:
            # Handle DataFrame, array, or CSV file input
            if isinstance(grain_data, pd.DataFrame):
                dfphi = grain_data.copy()
            elif isinstance(grain_data, (list, tuple, np.ndarray)):
                arr = np.asarray(grain_data)
                if arr.ndim == 2 and arr.shape[1] == 2:
                    dfphi = pd.DataFrame(arr, columns=["i(di)", "di[mm]"])
                elif arr.ndim == 2 and arr.shape[0] == 2:
                    dfphi = pd.DataFrame({"i(di)": arr[0, :], "di[mm]": arr[1, :]})
                else:
                    raise ValueError(
                        "grain_data array must be shape (n,2) or (2,n) with columns [i(di), di[mm]]"
                    )
            elif isinstance(grain_data, str):
                dfphi = pd.read_csv(grain_data)
            else:
                raise ValueError(
                    "grain_data must be float (D50), a 2D array/list with [i(di), di[mm]], "
                    "a DataFrame, or a path to a CSV file"
                )

            # Ensure numeric columns and sensible ordering
            if "i(di)" not in dfphi.columns or "di[mm]" not in dfphi.columns:
                raise ValueError("grain_data must provide columns 'i(di)' and 'di[mm]'")

            dfphi["i(di)"] = pd.to_numeric(dfphi["i(di)"], errors="coerce")
            dfphi["di[mm]"] = pd.to_numeric(dfphi["di[mm]"], errors="coerce")

            if dfphi["i(di)"].isnull().any() or dfphi["di[mm]"].isnull().any():
                raise ValueError("grain_data contains non-numeric values")

            # If i(di) appears to be percentages (0-100), convert to fractions (0-1)
            if dfphi["i(di)"].max() > 1.0:
                dfphi["i(di)"] = dfphi["i(di)"] / 100.0

            # Force cumulative behavior: sort by di and ensure monotonic cumulative values
            dfphi = dfphi.sort_values("di[mm]").reset_index(drop=True)
            # Clip cumulative to [0,1] and enforce non-decreasing
            dfphi["i(di)"] = dfphi["i(di)"].clip(0.0, 1.0)
            dfphi["i(di)"] = np.maximum.accumulate(dfphi["i(di)"])

            dfphi["di(Fehr) [mm]"] = dfphi["di[mm]"].interpolate()
            dfphi["Phi Scale"] = -np.log2(dfphi["di(Fehr) [mm]"])
            dfphi["Percent"] = (
                dfphi["i(di)"].diff().fillna(dfphi["i(di)"].iloc[0]) * 100
            )

            # Create bin edges: 18 classes need 19 edges
            # Edges at: -10, -9, -8, ..., 7, 8
            phi_bin_edges = np.concatenate([[-10], phi_class_centers + 0.5])

            # Bin the grain sizes into phi classes using pd.cut
            # This assigns each grain size to a phi class bin
            dfphi["Phi Class"] = pd.cut(
                dfphi["Phi Scale"],
                bins=phi_bin_edges,
                labels=phi_class_centers,
                include_lowest=True,
            )

            # Group by phi class and sum the percentages
            phi_percentages = dfphi.groupby("Phi Class", observed=False)[
                "Percent"
            ].sum()

            # Ensure ALL phi classes are present (fill missing with zeros)
            phi_percentages = phi_percentages.reindex(phi_class_centers, fill_value=0.0)

            # Normalize to ensure sum = 100%
            total = phi_percentages.sum()
            if total > 0:
                phi_percentages = phi_percentages / total * 100.0
            else:
                raise ValueError(
                    "Grain size distribution resulted in zero total percentage. "
                    "Check that grain_data covers a reasonable size range."
                )

            # Convert to fractions (0-1)
            phi_percentages = phi_percentages / 100.0

        # **VALIDATION: Verify length** (applies to both D50 and DataFrame/array paths)
        if len(phi_percentages) != expected_phi_length:
            raise ValueError(
                f"Grain size distribution produced {len(phi_percentages)} phi classes, "
                f"but expected {expected_phi_length} classes for range [-9.5, 7.5]. "
                f"This is an internal error - please report this issue."
            )

        # **VALIDATION: Verify sum** (applies to both D50 and DataFrame/array paths)
        total_fraction = phi_percentages.sum()
        if not np.isclose(total_fraction, 1.0, atol=0.01):
            raise ValueError(
                f"Phi percentages must sum to 1.0, got {total_fraction:.4f}. "
                f"Check grain size distribution normalization."
            )

        self.grain_size_data = dfphi
        self.phi_percentages = phi_percentages
        return self

    def export_scenarios_summary(
        self, output_path: str = None, format: str = "csv"
    ) -> DataFrame:
        """Export a comprehensive summary table of all scenarios with their parameters and indices.

        This method generates a user-friendly table containing, for each scenario:
        - Scenario metadata (name, description)
        - Scenario parameters (Qreq values, Qabs_max)
        - IHA/IARI indices (aggregated and by group)
        - Abstracted volumes (yearly totals, monthly averages)
        - Monthly released flows
        - Habitat indices (IH, ISH, ITH, HSD) if available
        - Annual sediment budget (total and per phi class) if available

        Parameters
        ----------
        output_path : str, optional
            Path where to save the export file. If None, only returns the DataFrame without saving.
        format : str, default='csv'
            Export format. Options: 'csv', 'excel'. Only used if output_path is provided.

        Returns
        -------
        DataFrame
            A pandas DataFrame containing the comprehensive summary of all scenarios.

        Raises
        ------
        ValueError
            If no scenarios have been added to the reach.
            If format is not 'csv' or 'excel'.

        Examples
        --------
        >>> reach = Reach("MyReach", dates, Qnat, Qabs_max)
        >>> # ... add scenarios and compute their metrics ...
        >>> df = reach.export_scenarios_summary("output.csv", format="csv")
        >>> df = reach.export_scenarios_summary("output.xlsx", format="excel")
        """
        if len(self.scenarios) == 0:
            raise ValueError("No scenarios have been added to this reach.")

        if format not in ["csv", "excel"]:
            raise ValueError("format must be 'csv' or 'excel'")

        # Collect data for all scenarios
        data_rows = []

        for scenario in self.scenarios:
            row = {
                "scenario_name": scenario.name,
                "scenario_description": scenario.description,
            }

            # Add scenario parameters
            row["Qabs_max"] = scenario.Qabs_max

            # Add scenario-specific parameters
            if hasattr(scenario, "Qreq_months"):
                row["scenario parameters"] = scenario.Qreq_months
                # for i, qr in enumerate(scenario.Qreq_months):
                #     row[f"Qreq_month_{i+1}"] = qr
            elif hasattr(scenario, "Qbase"):
                row["scenario parameters"] = [
                    scenario.Qbase,
                    scenario.c_Qin,
                    scenario.Qreq_min,
                    scenario.Qreq_max,
                ]

            # Add monthly released flows (average per month)
            if scenario.Qrel is not None:
                months = np.array([d.month for d in scenario.dates])
                for month in range(1, 13):
                    month_mask = months == month
                    if np.any(month_mask):
                        row[f"Qrel_mean_month_{month}_m3s"] = np.mean(
                            scenario.Qrel[month_mask]
                        )

            # Add volume statistics if available
            if hasattr(scenario, "yearly_abs_volumes") and hasattr(
                scenario, "yearly_nat_volumes"
            ):
                row["yearly_abs_volume_mean_m3"] = np.mean(scenario.yearly_abs_volumes)
                row["yearly_nat_volume_mean_m3"] = np.mean(scenario.yearly_nat_volumes)

                # Normalized abstracted volume (handle division by zero)
                # Set to 0 where natural volume is 0 to avoid division errors
                with np.errstate(divide="ignore", invalid="ignore"):
                    abs_norm = np.where(
                        scenario.yearly_nat_volumes != 0,
                        scenario.yearly_abs_volumes / scenario.yearly_nat_volumes,
                        0.0,
                    )
                row["abs_volume_normalized_mean"] = np.mean(abs_norm)

            # Add monthly abstracted volumes if available
            if hasattr(scenario, "monthly_abs_volumes"):
                for i, vol in enumerate(scenario.monthly_abs_volumes):
                    row[f"monthly_abs_volume_month_{i+1}_m3"] = vol

            # Add seasonal volumes if available
            if hasattr(scenario, "seasonal_abs_volumes"):
                for season, vol in scenario.seasonal_abs_volumes.items():
                    row[f"seasonal_abs_volume_{season}_m3"] = vol

            # Add cases duration if available
            if hasattr(scenario, "cases_duration"):
                row["case1_duration_fraction"] = scenario.cases_duration[0]
                row["case2_duration_fraction"] = scenario.cases_duration[1]
                row["case3_duration_fraction"] = scenario.cases_duration[2]

            # Add habitat indices if available
            ih_dict = getattr(scenario, "IH", {})
            if ih_dict:
                for species, ih_data in ih_dict.items():
                    row[f"IH_{species}"] = ih_data.get("IH", np.nan)
                    row[f"ISH_{species}"] = ih_data.get("ISH", np.nan)
                    row[f"ITH_{species}"] = ih_data.get("ITH", np.nan)
                    row[f"HSD_{species}"] = ih_data.get("HSD", np.nan)

            # Add IHA indices if available (IARI)
            if hasattr(scenario, "IARI"):
                iari_dict = scenario.IARI
                row["IARI_aggregated_mean"] = np.mean(iari_dict["aggregated"])
                for group_name, group_values in iari_dict["groups"].items():
                    row[f"IARI_{group_name}_mean"] = np.mean(group_values)

            # Add normalized IHA indices if available
            if hasattr(scenario, "normalized_IHA"):
                norm_iha_dict = scenario.normalized_IHA
                row["normalized_IHA_aggregated_mean"] = np.mean(
                    norm_iha_dict["aggregated"]
                )
                for group_name, group_values in norm_iha_dict["groups"].items():
                    row[f"normalized_IHA_{group_name}_mean"] = np.mean(group_values)

            # Add annual sediment budget if available
            if hasattr(scenario, "annual_sediment_budget"):
                budget = scenario.annual_sediment_budget
                # If it's a DataFrame, compute mean values
                if isinstance(budget, pd.DataFrame):
                    # Add mean total sediment budget
                    if "Qs_total" in budget.columns:
                        row["annual_sediment_budget_total_mean"] = budget[
                            "Qs_total"
                        ].mean()

                # If it's a dict, compute mean from yearly values
                elif isinstance(budget, dict):
                    # Structure is {year: {phi_class: value, ...}}
                    all_years = list(budget.keys())
                    first_year_data = budget[all_years[0]]
                    if "Qs_total" in first_year_data:
                        total_values = [budget[year]["Qs_total"] for year in all_years]
                        row["annual_sediment_budget_total_mean"] = np.mean(total_values)

            data_rows.append(row)

        # Create DataFrame
        df = pd.DataFrame(data_rows)

        # Save to file if output_path is provided
        if output_path is not None:
            if format == "csv":
                df.to_csv(output_path, index=False)
            elif format == "excel":
                df.to_excel(output_path, index=False, engine="openpyxl")

        return df
