"""
CanopyPhotosynthesisModel: integrate shadow geometry with Farquhar model
to compute vine-level photosynthesis from zone-level PAR distribution.
"""

from __future__ import annotations

import numpy as np
import pandas as pd

from config.settings import FRUITING_ZONE_INDEX
from src.farquhar_model import FarquharModel
from src.solar_geometry import ShadowModel


class CanopyPhotosynthesisModel:
    """Compute vine-level A by running Farquhar on each canopy zone."""

    def __init__(
        self,
        shadow_model: ShadowModel | None = None,
        farquhar_model: FarquharModel | None = None,
        lai: float = 2.5,
        shade_temp_offset: float = -1.5,
        diffuse_fraction: float = 0.15,
    ):
        self.shadow = shadow_model or ShadowModel()
        self.farquhar = farquhar_model or FarquharModel()
        self.lai = lai
        self.shade_temp_offset = shade_temp_offset
        self.diffuse_fraction = diffuse_fraction

        # Zone weights from LAI distribution (bottom to top)
        nv = self.shadow.n_vertical
        nh = self.shadow.n_horizontal
        # Distribute LAI weights across zones
        vert_weights = self.shadow.lai_weights  # shape (n_vertical,)
        # Each horizontal zone within a row gets equal share
        self._zone_weights = np.outer(vert_weights, np.ones(nh) / nh)
        # Normalize so total = 1
        self._zone_weights /= self._zone_weights.sum()

    def compute_vine_A(
        self,
        par: float,
        Tleaf: float,
        CO2: float,
        VPD: float,
        Tair: float,
        shadow_mask: np.ndarray,
        solar_elevation: float | None = None,
        solar_azimuth: float | None = None,
        tracker_tilt: float | None = None,
    ) -> dict:
        """
        Compute vine-level A for a single timestep.

        Returns dict with:
            A_vine: weighted vine-level A (umol CO2 m-2 s-1)
            A_zones: array of A per zone (n_vertical x n_horizontal)
            sunlit_fraction: fraction of zones in sun
            par_zones: PAR per zone
        """
        par_zones = self.shadow.compute_par_distribution(
            par, shadow_mask, self.diffuse_fraction,
            solar_elevation=solar_elevation, solar_azimuth=solar_azimuth,
            tracker_tilt=tracker_tilt,
        )
        A_zones = np.zeros_like(par_zones)

        for iz in range(self.shadow.n_vertical):
            for ix in range(self.shadow.n_horizontal):
                zone_par = par_zones[iz, ix]
                # Shaded zones are slightly cooler
                zone_tleaf = Tleaf + (self.shade_temp_offset if shadow_mask[iz, ix] else 0.0)
                zone_tair = Tair + (self.shade_temp_offset * 0.5 if shadow_mask[iz, ix] else 0.0)

                if zone_par > 0:
                    A_zones[iz, ix] = self.farquhar.calc_photosynthesis(
                        PAR=zone_par, Tleaf=zone_tleaf, CO2=CO2,
                        VPD=VPD, Tair=zone_tair,
                    )

        A_vine = float(np.sum(A_zones * self._zone_weights)) * self.lai
        sunlit_frac = self.shadow.sunlit_fraction(shadow_mask)

        # Extract fruiting zone (zone 1) and top canopy (zone 2) summaries
        fz = FRUITING_ZONE_INDEX  # default 1
        top = min(self.shadow.n_vertical - 1, 2)  # zone 2 = apical

        fruiting_zone_A = float(A_zones[fz, :].mean()) if A_zones.shape[0] > fz else 0.0
        fruiting_zone_par = float(par_zones[fz, :].mean()) if par_zones.shape[0] > fz else 0.0
        top_canopy_A = float(A_zones[top, :].mean()) if A_zones.shape[0] > top else 0.0
        top_canopy_par = float(par_zones[top, :].mean()) if par_zones.shape[0] > top else 0.0

        return {
            "A_vine": A_vine,
            "A_zones": A_zones,
            "sunlit_fraction": sunlit_frac,
            "par_zones": par_zones,
            "fruiting_zone_A": fruiting_zone_A,
            "fruiting_zone_par": fruiting_zone_par,
            "top_canopy_A": top_canopy_A,
            "top_canopy_par": top_canopy_par,
        }

    def compute_timeseries(
        self,
        df: pd.DataFrame,
        shadow_masks: np.ndarray,
        par_col: str = "Air1_PAR_ref",
        tleaf_col: str = "Air1_leafTemperature_ref",
        co2_col: str = "Air1_CO2_ref",
        vpd_col: str = "Air1_VPD_ref",
        tair_col: str = "Air1_airTemperature_ref",
    ) -> pd.DataFrame:
        """
        Compute vine-level A for each row in df using pre-computed shadow masks.
        shadow_masks: array of shape (len(df), n_vertical, n_horizontal).
        """
        records = []
        for i, (_, row) in enumerate(df.iterrows()):
            par = float(row[par_col]) if pd.notna(row[par_col]) else 0.0
            tleaf = float(row[tleaf_col]) if pd.notna(row[tleaf_col]) else 25.0
            co2 = float(row[co2_col]) if pd.notna(row[co2_col]) else 400.0
            vpd = float(row[vpd_col]) if pd.notna(row[vpd_col]) else 1.5
            tair = float(row[tair_col]) if pd.notna(row[tair_col]) else 25.0

            mask = shadow_masks[i]
            result = self.compute_vine_A(par, tleaf, co2, vpd, tair, mask)

            # Also compute reference (no panel = no shadow)
            no_shadow = np.zeros_like(mask, dtype=bool)
            ref_result = self.compute_vine_A(par, tleaf, co2, vpd, tair, no_shadow)

            records.append({
                "A_vine_panel": result["A_vine"],
                "A_vine_ref": ref_result["A_vine"],
                "sunlit_fraction": result["sunlit_fraction"],
                "par_mean_panel": result["par_zones"].mean(),
                "par_mean_ref": ref_result["par_zones"].mean(),
            })
        return pd.DataFrame(records, index=df.index)