import numpy as np
import pandas as pd

# -----------------------------
# PARAMETERS & SCENARIOS
# -----------------------------

years = np.arange(2025, 2041)

base_params = {
    "initial_fleet": 12000000,  # total vehicles
    "initial_nev_share": 0.02,  # 2% NEVs in 2025
    "scrappage_rate": 0.055,

    # Cost assumptions
    "ice_tco": 14000,  # total cost of ownership per year (R)
    "nev_tco": 11000,
    "fuel_efficiency": 7,  # L/100km
    "fuel_price": 24,  # R/L
    "elec_per_km": 0.18,  # kWh/km
    "electricity_price": 2.8,  # R/kWh

    # Behavioural effects
    "price_sensitivity": 0.45,
    "fuel_price_sensitivity": 0.25,
    "awareness_growth": 0.03,
    "charger_effect": 0.15,

    # Economic relationships
    "jobs_base": 500000,
    "jobs_elasticity": 0.6,
    "fuel_tax_rate": 3.93,
    "vat_rate": 0.15,

    # Social effects
    "gini_base": 0.63,
    "gini_elasticity": 0.02,
    "air_quality_base": 100,
    "air_quality_sensitivity": 25,
}

# Corrected scenarios dictionary
scenarios = {
    "BAU": {
        "subsidy": 0,
        "extra_awareness": 0,
        "fuel_price_multiplier": 1,
        "charger_growth": 0.05
    },
    "SUBSIDY": {
        "subsidy": 30000,
        "extra_awareness": 0.05,
        "fuel_price_multiplier": 1,
        "charger_growth": 0.12
    },
    "FUEL_SHOCK": {
        "subsidy": 0,
        "extra_awareness": 0.02,
        "fuel_price_multiplier": 1.4,
        "charger_growth": 0.07
    }
}

# -----------------------------
# Adoption rate function
# -----------------------------
def adoption_rate(params, scenario, year, awareness, chargers, fuel_price, ice_tco, nev_tco):
    cost_diff = (ice_tco - nev_tco) / ice_tco
    fuel_cost_effect = (fuel_price - params["fuel_price"]) / params["fuel_price"]

    rate = (
        0.02
        + params["price_sensitivity"] * cost_diff
        + params["fuel_price_sensitivity"] * fuel_cost_effect
        + params["awareness_growth"] * awareness
        + scenario["extra_awareness"]
        + params["charger_effect"] * chargers
    )

    return np.clip(rate, 0.01, 0.95)

# -----------------------------
# Main model function
# -----------------------------
def run_model(params, scenario):
    results = []

    fleet = params["initial_fleet"]
    nev_share = params["initial_nev_share"]
    awareness = 0.2
    chargers = 0.03  # proxy

    for year in years:
        fuel_price = params["fuel_price"] * scenario["fuel_price_multiplier"]
        ice_tco = params["ice_tco"]
        nev_tco = params["nev_tco"] - scenario["subsidy"] / 5  # amortized

        # ADOPTION RATE
        adoption = adoption_rate(params, scenario, year, awareness, chargers,
                                 fuel_price, ice_tco, nev_tco)

        # FLOWS
        scrappage = fleet * params["scrappage_rate"]
        new_sales = fleet * 0.07
        nev_sales = new_sales * adoption

        # STOCKS
        nev_fleet = fleet * nev_share + nev_sales - (fleet * nev_share * params["scrappage_rate"])
        fleet = fleet + new_sales - scrappage
        nev_share = nev_fleet / fleet

        # ENERGY USE
        ice_km = 14000
        nev_km = 14000
        fuel_consumption = (fleet * (1 - nev_share)) * (ice_km / 100) * params["fuel_efficiency"]
        electricity_consumption = (fleet * nev_share) * nev_km * params["elec_per_km"]

        # EMPLOYMENT
        jobs = (
            params["jobs_base"] * (1 - params["jobs_elasticity"] * nev_share)
            + 120000 * nev_share
        )

        # GOVERNMENT REVENUE
        fuel_tax = fuel_consumption * params["fuel_tax_rate"]
        vat = new_sales * ice_tco * params["vat_rate"]
        subsidy_cost = nev_sales * scenario["subsidy"]
        gov_balance = fuel_tax + vat - subsidy_cost

        # HOUSEHOLD COSTS
        avg_cost = (1 - nev_share) * ice_tco + nev_share * nev_tco

        # SOCIAL METRICS
        gini = params["gini_base"] - params["gini_elasticity"] * nev_share
        air_quality = params["air_quality_base"] - params["air_quality_sensitivity"] * nev_share

        # Update awareness and chargers
        awareness = min(1, awareness + params["awareness_growth"])
        chargers += scenario["charger_growth"]

        results.append([
            year, fleet, nev_share, new_sales, nev_sales, fuel_consumption,
            electricity_consumption, jobs, gov_balance, avg_cost, gini, air_quality
        ])

    cols = [
        "Year", "Fleet", "NEV_Share", "New_Sales", "NEV_Sales", "Fuel_Consumption",
        "Electricity_Consumption", "Jobs", "Gov_Balance", "Household_Cost",
        "Gini", "Air_Quality"
    ]
    return pd.DataFrame(results, columns=cols)

# -----------------------------
# Run scenarios
# -----------------------------
outputs = {}
for name, sc in scenarios.items():
    outputs[name] = run_model(base_params, sc)

# Check SUBSIDY scenario
outputs["SUBSIDY"].head()