import csv
import json
import math
from itertools import product
from pathlib import Path

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np

ROOT = Path(__file__).resolve().parents[1]
DATA = ROOT / 'data/cost-metrics-2026-09-17'
OUT = ROOT / 'outputs/cost-metrics-2026-09-17'
HOURS = 8760


def crf(rate, years):
    return 1 / years if rate == 0 else rate / (1 - (1 + rate) ** -years)


def financing_factor(rate, years):
    return 1 if years == 0 else sum((1 + rate) ** (years - k - .5) for k in range(years)) / years


def save_rows(name, rows):
    (DATA / (name + '.json')).write_text(json.dumps(rows, indent=2) + '\n')
    with (DATA / (name + '.csv')).open('w') as f:
        writer = csv.DictWriter(f, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)


def save_plot(fig, name, footer):
    fig.text(.015, .015, footer, fontsize=8, color='#374151', va='bottom')
    fig.tight_layout(rect=(0,.105,1,.98))
    fig.savefig(OUT / (name + '.png'), dpi=170, bbox_inches='tight', pad_inches=.18)
    fig.savefig(OUT / (name + '.svg'), bbox_inches='tight', pad_inches=.18)
    plt.close(fig)


def financial_cases():
    rows=[]
    for cap,r,build,life,cf in product([6000,10000,15000],[.03,.05,.07,.10],[0,5,10,15],[30,60,80],[.60,.75,.90,.95]):
        factor=financing_factor(r,build)
        annual=cap*factor*crf(r,life)+120
        rows.append(dict(overnight_chf_kw=cap,real_rate=r,construction_years=build,life_years=life,
                         capacity_factor=cf,finance_factor=factor,financed_chf_kw=cap*factor,
                         fixed_om_chf_kw_year=120,fuel_variable_waste_chf_mwh=15,
                         partial_generation_cost_chf_mwh=annual/(HOURS*cf/1000)+15))
    save_rows('finance-utilization-sensitivities',rows)
    delays=[]
    for r,delay in product([.03,.05,.07,.10],[0,2,5]):
        capital=10000*financing_factor(r,10)*(1+r)**delay
        delays.append(dict(overnight_chf_kw=10000,real_rate=r,planned_build_years=10,idle_delay_years=delay,
                           financed_chf_kw=capital,life_years=60,capacity_factor=.9,
                           partial_generation_cost_chf_mwh=(capital*crf(r,60)+120)/(8.76*.9)+15))
    save_rows('idle-delay-sensitivities',delays)
    fig,ax=plt.subplots(figsize=(10,5.5))
    rates=[.03,.05,.07,.10]
    matrix=np.array([[next(v['partial_generation_cost_chf_mwh'] for v in rows if v['overnight_chf_kw']==10000 and v['real_rate']==r and v['construction_years']==b and v['life_years']==60 and v['capacity_factor']==.9) for b in [0,5,10,15]] for r in rates])
    im=ax.imshow(matrix,cmap='YlOrRd',aspect='auto',vmin=0)
    for i,j in product(range(4),range(4)):
        ax.text(j,i,f'{matrix[i,j]:.0f}',ha='center',va='center',color='white' if matrix[i,j]>210 else '#111827',fontsize=14)
    ax.set_xticks(range(4),['0\ninstant benchmark','5','10','15'])
    ax.set_yticks(range(4),['3%','5%','7%','10%'])
    ax.set_xlabel('Construction years; expenditure spread evenly at annual midpoints')
    ax.set_ylabel('Real financing / discount rate')
    ax.set_title('Finance changes the cost of the same hypothetical plant',loc='left',fontweight='bold',pad=16)
    fig.colorbar(im,ax=ax,label='Partial generation cost, CHF/MWh')
    save_plot(fig,'01_financing','ILLUSTRATION, not a Swiss nuclear quote. CHF10,000/kW overnight; 60 operating years; 90% utilization.\nFOM CHF120/kW-year + fuel/variable/waste allowance CHF15/MWh. Excludes grids, backup, taxes and unresolved dismantling/fund costs.\nSource: prior-owner-model assumptions; midpoint cash-flow calculation. No future cost reduction or support.')
    return rows


def storage_cases():
    rows=[]
    for h,cycles,life,r,charge in product([2,4,8,24],[20,100,250,365],[10,15,20],[.03,.07,.10],[0,40,80]):
        annual_discharge=h*cycles
        minimum_hours=h*cycles*(1+1/.85)
        feasible=minimum_hours<=HOURS
        cap=600000+225000*h
        storage_only=cap*(crf(r,life)+.04)/annual_discharge
        rows.append(dict(power_mw=1,usable_ac_energy_mwh=h,duration_hours=h,full_cycles_year=cycles,
                         life_years=life,real_rate=r,round_trip_efficiency=.85,
                         power_cost_chf_kw=600,energy_cost_chf_kwh=225,capital_chf=cap,
                         annual_fom_including_augmentation_fraction=.04,charge_price_chf_mwh=charge,
                         annual_discharge_mwh=annual_discharge,minimum_charge_discharge_hours=minimum_hours,
                         chronologically_possible_at_same_charge_power=feasible,
                         storage_service_cost_chf_mwh=storage_only if feasible else None,
                         charging_cost_per_discharged_mwh=charge/.85 if feasible else None,
                         lcos_including_charging_chf_mwh=storage_only+charge/.85 if feasible else None))
    save_rows('storage-duration-utilization-sensitivities',rows)
    fig,ax=plt.subplots(figsize=(10,5.5))
    matrix=np.full((4,4),np.nan)
    for i,c in enumerate([20,100,250,365]):
        for j,h in enumerate([2,4,8,24]):
            v=next(x for x in rows if x['duration_hours']==h and x['full_cycles_year']==c and x['life_years']==15 and x['real_rate']==.07 and x['charge_price_chf_mwh']==40)
            if v['lcos_including_charging_chf_mwh'] is not None:
                matrix[i,j]=v['lcos_including_charging_chf_mwh']
    cmap=plt.get_cmap('YlOrRd').copy()
    cmap.set_bad('#e5e7eb')
    from matplotlib.colors import LogNorm
    im=ax.imshow(matrix,cmap=cmap,norm=LogNorm(vmin=np.nanmin(matrix),vmax=np.nanmax(matrix)),aspect='auto')
    for i,j in product(range(4),range(4)):
        value=matrix[i,j]
        label='Time limit\nexceeded' if np.isnan(value) else f'{value:,.0f}'
        ax.text(j,i,label,ha='center',va='center',fontsize=12,color='white' if value>1800 else '#111827')
    ax.set_xticks(range(4),['2 h','4 h','8 h','24 h'])
    ax.set_yticks(range(4),['20','100','250','365'])
    ax.set_xlabel('Usable AC discharge duration at 1 MW')
    ax.set_ylabel('Full discharge cycles per year')
    ax.set_title('Storage used rarely has a high cost per delivered MWh',loc='left',fontweight='bold',pad=16)
    fig.colorbar(im,ax=ax,label='LCOS including charging, CHF/MWh (log colour scale)')
    save_plot(fig,'02_storage_utilization','ASSUMPTION SENSITIVITY, not tenders. Power CHF600/kW + usable energy CHF225/kWh; 15 years; 7% real.\n85% round-trip efficiency; charge electricity CHF40/MWh; 4% annual FOM already includes augmentation. No added augmentation bill.\nSources: prior-owner-model; NLR ATB methodology. Grey cells require more than 8,760 charging + discharging hours.')
    return rows


def solar_cases(observations):
    pv=[r for r in observations if r['id'].startswith('pv-roof-')]
    rows=[]
    for p,r,life,yield_kwh in product(pv,[.03,.05,.07],[20,25,30],[850,1000,1150]):
        annual_om=30 if p['id'] in ['pv-roof-1','pv-roof-2','pv-roof-3'] else 20
        rows.append(dict(observation_id=p['id'],power_band_kwp=p['power_band_kwp'],capital_chf_kwp=p['value'],
                         real_rate=r,life_years=life,annual_yield_kwh_kwp=yield_kwh,
                         implied_capacity_factor=yield_kwh/HOURS,annual_om_including_inverter_chf_kwp=annual_om,
                         partial_generation_cost_chf_mwh=(p['value']*crf(r,life)+annual_om)/(yield_kwh/1000)))
    save_rows('pv-generation-cost-sensitivities',rows)
    fig,ax=plt.subplots(figsize=(10,5.5))
    x=np.arange(len(pv))
    med=np.array([v['value'] for v in pv]); low=np.array([v['q1'] for v in pv]); high=np.array([v['q3'] for v in pv])
    ax.errorbar(x,med,yerr=[med-low,high-med],fmt='o',color='#155e75',ecolor='#0891b2',capsize=10,markersize=9,lw=3)
    for j,v in enumerate(pv):
        ax.text(j,v['q3']+105,f"{v['value']:,}\nn = {v['sample_n']}",ha='center',fontsize=10)
    ax.set_xticks(x,[r['power_band_kwp'] for r in pv])
    ax.set_ylim(0,3700);ax.set_xlabel('Rooftop installation capacity, kWp');ax.set_ylabel('CHF/kWp, excluding VAT')
    ax.set_title('Swiss rooftop prices fall with installation size',loc='left',fontweight='bold',pad=15)
    ax.grid(axis='y',alpha=.2)
    save_plot(fig,'03_swiss_rooftop_prices','OBSERVATION: 2025 offers and invoices; points = median; whiskers = middle 50%, not confidence intervals.\nInstaller scope; excludes external engineering, roof reinforcement, grid strengthening and batteries. Largest class has only 7 records.\nSource: Planair for EnergieSchweiz/BFE, Preisbeobachtungsstudie 2025, 3 August 2026, Table 4 p.17.')
    capture=[r for r in observations if r['id'].startswith('capture-')]
    fig,ax=plt.subplots(figsize=(10,5.5))
    for j,tech in enumerate(['PV','hydro','wind','biomass']):
        vals=[next(v['value'] for v in capture if v['technology']==tech and v['reporting_period']==f'2026-Q{q}') for q in [1,2]]
        bars=ax.bar(np.arange(2)+(j-1.5)*.19,vals,.18,label=tech.capitalize(),color=['#d97706','#0284c7','#059669','#64748b'][j])
        ax.bar_label(bars,fmt='%.2f',fontsize=9,padding=3)
    ax.set_xticks([0,1],['2026 Q1','2026 Q2']);ax.set_ylim(0,145)
    ax.set_ylabel('Generation-weighted reference price, CHF/MWh');ax.legend(ncol=4,loc='upper center',frameon=False)
    ax.set_title('A generated MWh earns different value across production profiles',loc='left',fontweight='bold',pad=15)
    save_plot(fig,'04_capture_prices','OBSERVATION, not generation cost or profit. Day-ahead prices weighted by measured quarter-hour injections, VAT excluded.\nDifferent measured cohorts; prices do not include capacity/flexibility value, contracts, support or retail self-consumption.\nSource: BFE Referenz-Marktpreise, 14 July 2026, pp.2–5. Two quarters do not establish a long-term causal forecast.')


def toy_mix():
    rte=.85
    storage_power=1/rte
    storage_energy=12
    storage_annual=(600000*storage_power+225000*storage_energy)*(crf(.07,15)+.04)
    capacity=1+1/rte
    rows=[]
    for label,a,b,p,e,cost in [
        ('A only + storage',capacity,0,storage_power,storage_energy,capacity*100000+storage_annual),
        ('B only + storage',0,capacity,storage_power,storage_energy,capacity*120000+storage_annual),
        ('A + complementary B',1,1,0,0,220000),
        ('Ideal dispatchable C',0,0,0,0,250000),
        ('A + correlated B',capacity,0,storage_power,storage_energy,capacity*100000+storage_annual)]:
        rows.append(dict(case=label,a_mw=a,b_mw=b,c_mw=1 if label=='Ideal dispatchable C' else 0,storage_power_mw=p,usable_storage_mwh=e,
                         annual_cost_chf=cost,served_demand_mwh=HOURS,cost_chf_mwh=cost/HOURS))
    save_rows('toy-complementarity',rows)
    fig,ax=plt.subplots(figsize=(10,5.5))
    bars=ax.barh([v['case'] for v in rows],[v['cost_chf_mwh'] for v in rows],color=['#64748b','#64748b','#0891b2','#7c3aed','#64748b'])
    ax.invert_yaxis();ax.bar_label(bars,fmt='%.1f',padding=4);ax.set_xlim(0,max(x['cost_chf_mwh'] for x in rows)*1.16)
    ax.set_xlabel('Illustrative cost per served MWh, CHF');ax.set_title('The service requirement and the mix change the answer',loc='left',fontweight='bold',pad=15)
    save_plot(fig,'05_toy_complementarity','TEACHING EXAMPLE — no Swiss weather or technology forecast. Demand 1 MW; repeating 12-hour A / 12-hour B availability.\nA costs CHF100,000/MW-year; B120,000; ideal always-available C250,000. Battery 85% efficient; cyclic state; prior-model costs.\nPerfect complementarity is invented; correlated B matches A and brings no benefit here. No outages, networks or reserves in ANY case.')
    assert math.isclose(capacity*12,12+12/rte)
    assert rows[2]['cost_chf_mwh']<rows[3]['cost_chf_mwh']<rows[0]['cost_chf_mwh']
    assert rows[0]['annual_cost_chf']==rows[4]['annual_cost_chf']


def common_horizon():
    rows=[]
    for horizon,life,r,residual_mode in product([20,40,60],[15,25,60],[.03,.07],['zero','straight_line_remaining_life']):
        times=list(range(0,horizon,life))
        cost=sum(1000/(1+r)**t for t in times)
        remaining=max(0,times[-1]+life-horizon)
        residual=1000*remaining/life if residual_mode!='zero' else 0
        net=cost-residual/(1+r)**horizon
        rows.append(dict(horizon_years=horizon,asset_life_years=life,real_rate=r,residual_mode=residual_mode,
                         initial_and_replacement_cost_each_chf=1000,number_of_purchases=len(times),
                         terminal_residual_chf=residual,net_present_capital_chf=net,
                         equivalent_annual_capital_chf=net*crf(r,horizon)))
    save_rows('common-horizon-residual-sensitivities',rows)
    complete=next(x for x in rows if x['horizon_years']==60 and x['asset_life_years']==15 and x['real_rate']==.07 and x['residual_mode']=='zero')
    assert math.isclose(complete['equivalent_annual_capital_chf'],1000*crf(.07,15),rel_tol=1e-12)
    examples=[dict(id='charging-once',generation_mwh=100,delivered_mwh=85,generation_cost_chf_mwh=50,
                   storage_service_cost_chf=1000,correct_total_chf=6000,wrong_double_charge_total_chf=11000),
              dict(id='grant-transfer',gross_project_resource_cost_chf=100,grant_chf=60,
                   private_initial_payment_chf=40,public_payment_chf=60,national_resource_cost_chf=100)]
    (DATA/'double-counting-checks.json').write_text(json.dumps(examples,indent=2)+'\n')


def main():
    OUT.mkdir(parents=True,exist_ok=True)
    plt.rcParams.update({'font.family':'DejaVu Sans','axes.spines.top':False,'axes.spines.right':False})
    obs=json.loads((DATA/'cost-observations.json').read_text())['observations']
    financial_cases()
    storage_cases()
    solar_cases(obs)
    toy_mix()
    common_horizon()
    assert math.isclose(crf(0,20),.05)
    assert financing_factor(.07,15)>financing_factor(.07,10)>1
    print('Generated five visual examples and seven sensitivity datasets.')


if __name__=='__main__':
    main()