import pandas as pd
from typing import List, Dict, Any, Set
import backend.app.assumptions as assumptions
from backend.Property import Property
from backend.app.plan.schemas import PlanTriggerRequest
from backend.app.utils import epc_to_sap_lower_bound
from recommendations.optimiser.CostOptimiser import CostOptimiser


def prepare_input_measures(
    property_recommendations, goal, needs_ventilation, funding=False,
    property_eco_packages=None
):
    """
    Prepares a nested list of measure options for optimisation.

    Each sublist represents all available variants of a single measure type (e.g. all solar PV options).
    Within each sublist, each measure is represented as a dictionary containing:
        - id: unique recommendation identifier
        - cost: total cost of the measure (including ventilation if bundled)
        - gain: the relevant gain metric based on the selected goal
        - type: the measure type, optionally combined with ventilation (e.g. "wall_insulation+mechanical_ventilation")

    Ventilation bundling:
        - If a property needs ventilation, and a measure type requires it (as defined in
        assumptions.measures_needing_ventilation),
          the ventilation cost and gain are added to that measure’s values.

    Filtering:
        - Measures with negative `energy_cost_savings` are excluded.
        - Solar PV options with batteries are excluded (currently handled by a placeholder bitwise NOT).

    Parameters
    ----------
    property_recommendations : list[list[dict]]
        Nested list of recommendations for a property. Each inner list represents variations of the same measure type.
    goal : str
        Optimisation goal, one of: "Increasing EPC", "Energy Savings", "Reducing CO2 emissions".
    needs_ventilation : bool
        Whether the property requires mechanical ventilation to accompany certain measures.
    funding: bool, optional
        If true, the function will include the innovation uplift in the total cost calculation. If false, this is
        excluded, since innovation uplift cannot be claimed where funding is not available.
    property_eco_packages: dict, optional
        Eco package data for the property, if available. If a measure has been specified as part of an eco package
        (e.g. HHRSH) this function will include that measure in the optimisation, even if it has negative cost savings.

    Returns
    -------
    list[list[dict]]
        Nested list of prepared measure options, ready for input into the optimiser.
    """

    goal_map = {
        "Increasing EPC": "sap_points",
        "Energy Savings": "kwh_savings",
        "Reducing CO2 emissions": "co2_equivalent_savings",
    }

    goal_key = goal_map[goal]
    if not goal_key:
        raise NotImplementedError("Not implemented this gain type - investigate me")

    ventilation_recommendation = next(
        (measure[0] for measure in property_recommendations if measure[0]["type"] == "mechanical_ventilation"),
        {}
    )

    eco_measures = property_eco_packages[0] if property_eco_packages else []

    input_measures = []
    for recs in property_recommendations:

        # Skip ventilation as a standalone optimisation option (it will be bundled)
        if needs_ventilation and recs[0]["type"] == "mechanical_ventilation":
            continue

        # Filter out solar PV with batteries
        # if recs[0]["type"] == "solar_pv":
        #     recs = [r for r in recs if ~r["has_battery"]]

        # Only include measures with non-negative cost savings - we allow for a minor negative impact
        if eco_measures:
            recs_to_append = [
                rec for rec in recs if (rec["energy_cost_savings"] >= -10) or (rec["measure_type"] in eco_measures)
            ]
        else:
            recs_to_append = [
                rec for rec in recs if (rec["energy_cost_savings"] >= -10)
            ]
        if not recs_to_append:
            continue

        # Build enriched measure data
        to_append = []
        for rec in recs:

            raw_cost = rec["total"]

            if funding:
                total = (
                    rec["total"] - rec["innovation_uplift"] + ventilation_recommendation["total"]
                    if rec["measure_type"] in assumptions.measures_needing_ventilation and needs_ventilation
                    else rec["total"] - rec["innovation_uplift"]
                )
            else:
                total = (
                    rec["total"] + ventilation_recommendation["total"]
                    if rec["measure_type"] in assumptions.measures_needing_ventilation and needs_ventilation
                    else rec["total"]
                )

            # If the innovation uplift being removed make this negative, we keep the total so we can re-engineer
            # the original cost
            non_negative_total = 0 if total < 0 else total
            gain = (
                rec[goal_key] + ventilation_recommendation[goal_key]
                if rec["measure_type"] in assumptions.measures_needing_ventilation and needs_ventilation
                else rec[goal_key]
            )
            rec_type = (
                f"{rec['measure_type']}+{ventilation_recommendation['measure_type']}"
                if rec["measure_type"] in assumptions.measures_needing_ventilation and needs_ventilation
                else rec["measure_type"]
            )

            array_size = 0
            if rec["measure_type"] == "solar_pv":
                # Grab the parts
                solar_part = next(
                    (part for part in rec["parts"] if part["type"] == "solar_pv"),
                )
                array_size = solar_part["size"]

            # We also include the innovation uplift
            to_append.append(
                {
                    "id": rec["recommendation_id"],
                    "cost": non_negative_total,
                    "gain": gain, "type": rec_type,
                    "innovation_uplift": rec["innovation_uplift"] if funding else 0,
                    "cost_minus_uplift": total,
                    "raw_cost": raw_cost,
                    "partial_project_funding": rec["partial_project_funding"],
                    "partial_project_score": rec["partial_project_score"],
                    "uplift_project_score": rec["uplift_project_score"],
                    "already_installed": rec.get("already_installed", False),
                    "has_battery": rec.get("has_battery", False),
                    "array_size": array_size,
                }
            )

        input_measures.append(to_append)

    return input_measures


def calculate_fixed_gain(property_required_measures, recommendations, p, needs_ventilation):
    """
    Calculates the total "fixed gain" from required measures for a property.

    Required measures are applied regardless of optimisation. This function:
      - Finds the maximum SAP points for each required measure type.
      - Sums those max SAP values into a fixed gain total.
      - Adds the SAP points for mechanical ventilation if:
            * The property needs ventilation, and
            * At least one required measure needs ventilation.

    Parameters
    ----------
    property_required_measures : list[list[dict]]
        Nested list of required measures for the property.
    recommendations : dict
        All recommendations for all properties, keyed by property id.
    p : object
        Property object (must have .id).
    needs_ventilation : bool
        Whether ventilation should be bundled with certain measures.

    Returns
    -------
    float
        Total fixed SAP gain from required measures (and ventilation, if applicable).
    """
    if not property_required_measures:
        return 0

    sap_by_type = [
        {"type": rec["type"], "sap_points": rec["sap_points"]}
        for recs in property_required_measures for rec in recs
    ]

    max_per_type = pd.DataFrame(sap_by_type).groupby("type")["sap_points"].max().to_dict()
    fixed_gain = sum(max_per_type.values())

    required_types = {rec["type"] for rec in sap_by_type}
    if needs_ventilation and any(r in required_types for r in assumptions.measures_needing_ventilation):
        fixed_gain += next(
            (r[0]["sap_points"] for r in recommendations[p.id] if r[0]["type"] == "mechanical_ventilation"),
            0
        )

    return fixed_gain


def calculate_gain(
    body: PlanTriggerRequest,
    p: Property,
    fixed_gain: float,
    already_installed_gain: float = 0,
) -> float | None:
    """
    Calculates the target gain value for optimisation based on the goal.

    - For "Increasing EPC": Computes the SAP gain needed to reach the target EPC,
      applies a slack adjustment (via CostOptimiser), and subtracts fixed gains from required measures.
    - For "Energy Savings" or "Reducing CO2 emissions": Returns None,
      which signals the optimiser to simply maximise gain under a budget.

    Parameters
    ----------
    body : object
        Request body object containing optimisation settings (goal, goal_value, simulate_sap_10, etc.)
    p : object
        Property object with EPC data (must have p.data["current-energy-efficiency"]).
    fixed_gain : float
        Total fixed gain from required measures (returned by calculate_fixed_gain).
    already_installed_gain: float, optional

    Returns
    -------
    float or None
        Required SAP gain for EPC, or None for non-EPC goals.
    """
    if body.goal == "Increasing EPC":
        current_sap = p.epc_record.current_energy_efficiency + already_installed_gain
        target_sap = epc_to_sap_lower_bound(body.goal_value)

        if target_sap <= current_sap:
            # We've already met or exceeded the target EPC
            return 0

        gain = CostOptimiser.calculate_sap_gain_with_slack(
            target_sap - current_sap
        ) - fixed_gain
        if body.simulate_sap_10:
            gain += 3
        return max(gain, 0)
    elif body.goal in ["Energy Savings", "Reducing CO2 emissions"]:
        return None
    else:
        raise NotImplementedError(f"Goal {body.goal} is not supported")


def add_required_measures(property_id, property_required_measures, recommendations, selected):
    """
    Ensures the cheapest variant of each required measure is added to the selected recommendations.

    For each required measure type, this function:
      - Finds the lowest-cost variant not already selected.
      - Adds it to the selected recommendation IDs.
      - Returns a flattened list of all selected measure details for final output.

    Parameters
    ----------
    property_id : int
        Unique identifier for the property.
    property_required_measures : list[list[dict]]
        Nested list of required measures for the property.
    recommendations : dict
        All recommendations for all properties, keyed by property id.
    selected : set
        Set of already selected recommendation IDs from the optimiser.

    Returns
    -------
    list[dict]
        Flat list of selected measure details, each containing:
        {"id", "cost", "gain", "type"}
    """
    for recs in property_required_measures:
        cheapest = min(
            (rec for rec in recs if rec["recommendation_id"] not in selected),
            key=lambda rec: rec["total"],
        )
        selected.add(cheapest["recommendation_id"])

    return [
        {"id": rec["recommendation_id"], "cost": rec["total"], "gain": rec["sap_points"], "type": rec["type"]}
        for recs in recommendations[property_id] for rec in recs
        if rec["recommendation_id"] in selected
    ]


def add_best_practice_measures(
    property_id: int,
    solution: List[Dict[str, Any]],
    recommendations: Dict[int, List[List[Dict[str, Any]]]],
    selected: Set[str],
):
    """
    Ensures best-practice measures like ventilation and trickle vents are included
    in the selected recommendations when appropriate.

    Rules:
      - If a measure requiring ventilation is selected AND ventilation is not already present,
        add the corresponding mechanical ventilation recommendation.
      - Always add trickle vents if they exist in the recommendations.

    Parameters
    ----------
    property_id : int
        The unique identifier for the property.
    solution : list[dict]
        The current list of selected measures (each containing id, type, gain, cost).
    recommendations : dict
        All recommendations for all properties, keyed by property id.
    selected : set
        Set of already selected recommendation IDs.

    Returns
    -------
    set
        Updated set of selected recommendation IDs, including ventilation and trickle vents if applicable.
    """
    # Check if any selected measure requires ventilation
    ventilation_selected = [r for r in solution if "+mechanical_ventilation" in r["type"]]

    # If ventilation has been selected, or one of the measures needs ventilation, we need to ensure ventilation is
    # included
    measures_selected_needing_ventilation = any(
        x in [r["type"] for r in solution] for x in assumptions.measures_needing_ventilation
    )

    if measures_selected_needing_ventilation or len(ventilation_selected) > 0:
        ventilation_rec = next(
            (r[0] for r in recommendations[property_id] if r[0]["type"] == "mechanical_ventilation"),
            None
        )
        if ventilation_rec:
            selected.add(ventilation_rec["recommendation_id"])

    # Always add trickle vents if available
    trickle_vents_rec = next(
        (r[0] for r in recommendations[property_id] if r[0]["type"] == "trickle_vents"),
        None
    )
    if trickle_vents_rec:
        selected.add(trickle_vents_rec["recommendation_id"])

    return selected


def flatten_recommendations_with_defaults(property_id, recommendations, selected, battery_sap_score=0):
    """
    Flattens nested recommendation lists for a property and marks which
    recommendations were selected.

    Each recommendation dict is copied and an extra key `default` is added:
        - True if the recommendation ID is in `selected`
        - False otherwise

    Parameters
    ----------
    property_id : int
        The unique identifier for the property.
    recommendations : dict
        All recommendations for all properties, keyed by property id.
        Each value is a list of lists (grouped by measure type).
    selected : set
        Set of selected recommendation IDs.
    battery_sap_score: int, optional
        SAP score uplift from battery storage, if applicable.

    Returns
    -------
    list[dict]
        A flattened list of recommendation dicts for the given property,
        each with an added `default` field.
    """

    final_recommendations = []
    for recommendations_by_type in recommendations[property_id]:
        recs_by_type = []
        for rec in recommendations_by_type:
            rec_copy = {**rec, "default": rec["recommendation_id"] in selected}
            if rec_copy.get("has_battery", False):
                rec_copy["sap_points"] += battery_sap_score
            recs_by_type.append(rec_copy)

        final_recommendations.append(recs_by_type)

    # Flatten the nested list of lists into a single list
    return [rec for recommendations_by_type in final_recommendations for rec in recommendations_by_type]


def check_needs_ventilation(
    property_measure_types: Set[str],
    measures_needing_ventilation: List[str],
    property_already_has_ventilation: bool,
    ventilation_in_included_measures: bool
) -> bool:
    """
    Function to check if we need to include ventilation based on the measures selected and the property
    features
    :param property_measure_types: The set of measure types recommended for the property
    :param measures_needing_ventilation: The set of measure types that require ventilation
    :param property_already_has_ventilation: Whether the property currently has ventilation
    :param ventilation_in_included_measures: Whether ventilation is already included in the recommended
    measures
    :return: Boolean indicating whether ventilation needs to be included in the recommendations

    # TODO - none of the inputs of this function are well structured and so this is quite brittle - we should
             consider refactoring to make this more robust
    """

    needs_ventilation = any(
        x in property_measure_types for x in measures_needing_ventilation
    )

    return needs_ventilation and not property_already_has_ventilation and ventilation_in_included_measures
