import argparse
import copy
import json
from pathlib import Path

import numpy as np
import pandas as pd

from feasibility_model.dispatch import solve
from hourly_model.dispatch import hydro_year_mask, supply_profile
from hourly_model.load import demand_profile, read_weather

ROOT = Path(__file__).resolve().parents[1]
DATA = ROOT / 'data/system-feasibility-2026-09-10'


def write_json(path, value):
    path.write_text(json.dumps(value, indent=2, ensure_ascii=False) + '\n')


def options():
    stores = [
        dict(id='nant_de_drance', pump_gw=.9, turbine_gw=.9, deliverable_energy_gwh=20, roundtrip_efficiency=.8),
        dict(id='limmern', pump_gw=1, turbine_gw=1, deliverable_energy_gwh=23e6 * 1000 * 9.81 * 630 * .9 / 3.6e12, roundtrip_efficiency=.8),
        dict(id='hongrin_leman', pump_gw=.420, turbine_gw=.420, deliverable_energy_gwh=100, roundtrip_efficiency=.8),
        dict(id='grimsel2', pump_gw=.392, turbine_gw=.388, deliverable_energy_gwh=57e6 * 1000 * 9.81 * 400 * .9 / 3.6e12, roundtrip_efficiency=.8)]
    return dict(pumped_stores=stores, natural_capacity_gwh=8722 - sum(s['deliverable_energy_gwh'] for s in stores), natural_turbine_gw=12.156 * .98 - sum(s['turbine_gw'] for s in stores), battery_rte=.85, battery_max_gw=20, battery_max_gwh=10000, export_gw=8, battery_power_chf_per_kw=600, battery_energy_chf_per_kwh=225)


def case_definitions():
    cases = []
    for units, solar in [(8, 26), (6, 40), (3, 65), (0, 90)]:
        for year in [2012, 2018, 2022]:
            cases.append(dict(id=f'n{units}-pv{solar}-{year}-split', units=units, solar=solar, year=year, variant='split'))
    for variant in ['battery_none', 'battery_4h', 'road_flex', 'thermal_flex', 'hydro_75pct', 'hydro_50pct', 'compound', 'hydro_expansion', 'pump_residual', 'cold_heat', 'connected_2gw', 'hydro_power_75pct']:
        cases.append(dict(id=f'n8-pv26-2012-{variant}', units=8, solar=26, year=2012, variant=variant))
    for year in [2021, 2023, 2024]:
        cases.append(dict(id=f'n8-pv26-{year}-split', units=8, solar=26, year=year, variant='split'))
    for power in [1, 2, 4, 8, 10]:
        cases.append(dict(id=f'n8-pv26-2012-power{power}', units=8, solar=26, year=2012, variant='fixed_power', power=power))
    for units, solar in [(3, 65), (0, 90)]:
        for year in [2022, 2024]:
            cases.append(dict(id=f'n{units}-pv{solar}-{year}-wind2-flex', units=units, solar=solar, year=year, variant='wind_flex', wind_gw=2))
    return cases


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--case', default='all')
    args = parser.parse_args()
    config = json.loads((ROOT / 'hourly_model/config.json').read_text())
    weather = read_weather()
    coefficients = json.loads((ROOT / 'data/hourly-model-v1/load-fit.json').read_text())['all_years_shape_coefficients']
    demand, demand_audit = demand_profile(weather, coefficients, config)
    write_json(DATA / 'demand-reuse-audit.json', demand_audit)
    results = []
    for case in case_definitions():
        if args.case != 'all' and case['id'] != args.case:
            continue
        folder = DATA / 'cases' / case['id']
        if (folder / 'result.json').exists():
            results.append(json.loads((folder / 'result.json').read_text()))
            continue
        folder.mkdir(parents=True, exist_ok=True)
        p = options()
        supply = supply_profile(weather, config, dict(new_nuclear_units=case['units'], solar_gwp=case['solar']), outage=True)
        mask = hydro_year_mask(weather.index, case['year'])
        d, s = demand.loc[mask].copy(), supply.loc[mask].copy()
        variant = case['variant']
        if variant == 'battery_none':
            p.update(battery_fixed_gw=0, battery_fixed_gwh=0)
        if variant == 'battery_4h':
            p['battery_duration_hours'] = 4
        if variant == 'fixed_power':
            p['battery_fixed_gw'] = case['power']
        if variant == 'road_flex':
            p['road_flexible_fraction'] = .5
        if variant == 'thermal_flex':
            p.update(road_flexible_fraction=.5, thermal_shift_gw=1, thermal_shift_hours=6)
        if variant in ['hydro_75pct', 'hydro_50pct']:
            p['natural_capacity_gwh'] *= .75 if variant == 'hydro_75pct' else .5
        if variant == 'compound':
            s.reservoir_inflow_gw *= .8
            s.river_gw *= .8
            s.solar_gw *= .9
            local = d.index.tz_convert('Europe/Zurich')
            outage_mask = (local.month <= 3)
            s.loc[outage_mask, 'nuclear_gw'] = np.maximum(s.loc[outage_mask, 'nuclear_gw'] - 1.1 * .98, 0)
            p['natural_capacity_gwh'] *= .75
            p['natural_initial_fraction'] = .7
        if variant == 'cold_heat':
            local_weather = weather.loc[mask]
            weights = json.loads((ROOT / 'data/hourly-model-v1/weather-audit.json').read_text())['station_weights']
            temperature = sum(weight * local_weather['temperature_' + station + '_c'] for station, weight in weights.items())
            extra = d.space_heat * np.clip((0 - temperature) / 10, 0, 1) * .5
            d['space_heat'] += extra
            d['final_demand_gw'] += extra
            d['bus_demand_gw'] += extra / .94
        if variant == 'hydro_expansion':
            p['natural_capacity_gwh'] += 240 + 65 + 60
        if variant == 'pump_residual':
            p['pumped_stores'].append(dict(id='unmapped_residual', pump_gw=1, turbine_gw=1, deliverable_energy_gwh=40, roundtrip_efficiency=.8))
            p['natural_capacity_gwh'] -= 40
            p['natural_turbine_gw'] -= 1
        if variant == 'connected_2gw':
            p['import_gw'] = 2
        if variant == 'hydro_power_75pct':
            p['natural_turbine_gw'] *= .75
            for store in p['pumped_stores']:
                store['pump_gw'] *= .75
                store['turbine_gw'] *= .75
        if variant == 'wind_flex':
            raw = json.loads((ROOT / f'backups/system-feasibility-2026-09-10/sources/ch-public-power-{case["year"]}.json').read_text())
            wind_index = pd.to_datetime(raw['unix_seconds'], unit='s', utc=True)
            wind = np.asarray(next(t['data'] for t in raw['production_types'] if t['name'] == 'Wind onshore'), dtype=float)
            if not wind_index.equals(d.index) or not np.isfinite(wind).all() or wind.min() < 0:
                raise ValueError('Wind source calendar mismatch')
            profile = np.minimum(2500 * wind / wind.sum(), 1)
            s['wind_gw'] = case['wind_gw'] * profile
            s['other_gw'] += s.wind_gw
            p.update(road_flexible_fraction=.5, thermal_shift_gw=1, thermal_shift_hours=6)
            p['wind_profile_assumption'] = dict(source_id=f'ch-public-power-{case["year"]}', target_full_load_hours=2500, realized_full_load_hours=float(profile.sum()), status='Current small-fleet hourly shape scaled to assumed 2 GW future capacity and clipped at nameplate; not a validated future site portfolio. No baseline wind was credited.')
        s['fixed_generation_gw'] = s.solar_gw + s.nuclear_gw + s.river_gw + s.other_gw
        frame, result = solve(d.bus_demand_gw, s.fixed_generation_gw, s.reservoir_inflow_gw, p, road=d.road / .94, heat=d.space_heat / .94, index=d.index)
        result.update(case=case, options=p, interpretation='Conditional national-bus screening with finite isolated pumped stores; residual natural hydro remains aggregate. Capacities minimize an assumed battery initial-capital proxy for this fixed generation portfolio and weather case. No merchant-profit or buildability conclusion.')
        write_json(folder / 'result.json', result)
        pd.concat([d, s, frame], axis=1).to_csv(folder / 'hourly.csv.gz', index_label='interval_start_utc')
        results.append(result)
        print(case['id'], 'GW', round(result['battery_gw'], 3), 'GWh', round(result['battery_deliverable_gwh'], 2), 'shortage', round(result['unserved_gwh'], 4), 'seconds', round(result['solve_seconds'], 1), flush=True)
    write_json(DATA / ('case-index.json' if args.case == 'all' else 'case-index-' + args.case + '.json'), results)


if __name__ == '__main__':
    main()