import json
from time import perf_counter

import numpy as np
import pandas as pd
from scipy.optimize import linprog
from scipy.sparse import coo_matrix

from hourly_model.load import local_index
from hourly_model.prepare import OUT, ZONE

VARIABLES = ['hydro_release', 'hydro_charge', 'hydro_stock', 'hydro_spill', 'battery_charge', 'battery_discharge', 'battery_stock', 'imports', 'exports', 'unserved', 'curtailment']


def hydro_year_mask(index, year):
    start = pd.Timestamp(f'{year - 1}-10-01', tz=ZONE).tz_convert('UTC')
    end = pd.Timestamp(f'{year}-10-01', tz=ZONE).tz_convert('UTC')
    return (index >= start) & (index < end)


def supply_profile(weather, config, portfolio, outage=False):
    p = config['supply']
    local = local_index(weather.index)
    weights = json.load(open(OUT / 'weather-audit.json'))['station_weights']
    radiation = sum(weights[s] * weather['radiation_' + s + '_w_m2'].to_numpy() for s in weights) / 1000
    reference = local.year == 2024
    low, high = 0., 10.
    for _ in range(60):
        midpoint = (low + high) / 2
        if np.minimum(radiation[reference] * midpoint, 1).sum() < p['solar_reference_kwh_per_kwp']:
            low = midpoint
        else:
            high = midpoint
    pv_per_gwp = np.minimum(radiation * (low + high) / 2, 1)
    hydro = json.load(open(OUT / 'hydro-monthly.json'))['rows']
    lookup = {(r['year'], r['month']): r for r in hydro}
    river = np.empty(len(weather))
    inflow = river.copy()
    for year, month in sorted(set(zip(local.year, local.month))):
        mask = (local.year == year) & (local.month == month)
        start = pd.Timestamp(year=int(year), month=int(month), day=1, tz=ZONE)
        end = start + pd.offsets.MonthBegin(1)
        hours = (end.tz_convert('UTC') - start.tz_convert('UTC')).total_seconds() / 3600
        row = lookup[(year, month)]
        river[mask] = row['run_of_river_gwh'] / hours * p['hydro_bus_delivery_fraction']
        inflow[mask] = row['derived_reservoir_inflow_gwh'] / hours * p['hydro_bus_delivery_fraction']
    river_clipped = np.minimum(river, p['run_of_river_gw_gross'] * p['hydro_bus_delivery_fraction'])
    nuclear = np.zeros(len(weather))
    n = portfolio['new_nuclear_units']
    for unit in range(n):
        maintenance_start_day = 91 + (unit % 6) * p['nuclear_maintenance_days_per_unit']
        maintenance = (local.dayofyear >= maintenance_start_day) & (local.dayofyear < maintenance_start_day + p['nuclear_maintenance_days_per_unit'])
        available = ~maintenance
        if outage and unit == 0:
            available &= ~((local.dayofyear >= 1) & (local.dayofyear <= p['nuclear_forced_outage_days']))
        nuclear += available * p['nuclear_unit_net_gw'] * p['nuclear_other_availability']
    result = pd.DataFrame({'solar_gw': pv_per_gwp * portfolio['solar_gwp'], 'nuclear_gw': nuclear, 'river_gw': river_clipped, 'river_clipping_gw': river - river_clipped, 'reservoir_inflow_gw': inflow, 'other_gw': p['other_generation_gw']}, index=weather.index)
    result['fixed_generation_gw'] = result.solar_gw + result.nuclear_gw + result.river_gw + result.other_gw
    return result


def solve_dispatch(demand, generation, inflow, config, import_limit_gw):
    started = perf_counter()
    demand, generation, inflow = map(np.asarray, [demand, generation, inflow])
    n = len(demand)
    if n == 0 or any(len(a) != n or not np.isfinite(a).all() or a.min() < 0 for a in [demand, generation, inflow]):
        raise ValueError('Invalid dispatch input')
    p = config['supply']
    eta_h = p['hydro_pump_roundtrip_efficiency']
    eta_b = np.sqrt(p['battery_roundtrip_efficiency'])
    capacity_h = p['combined_reservoir_gwh_gross'] * p['hydro_bus_delivery_fraction']
    power_h = p['reservoir_turbine_gw_gross'] * p['hydro_bus_delivery_fraction']
    indexes = {name: i * n + np.arange(n) for i, name in enumerate(VARIABLES)}
    rr, cc, vv = [], [], []
    def add(rows, columns, values):
        rr.extend(rows)
        cc.extend(columns)
        vv.extend(np.broadcast_to(values, n))
    hours = np.arange(n)
    for name, sign in [('hydro_release', 1), ('hydro_charge', -1), ('battery_charge', -1), ('battery_discharge', 1), ('imports', 1), ('exports', -1), ('unserved', 1), ('curtailment', -1)]:
        add(hours, indexes[name], sign)
    add(n + hours, indexes['hydro_stock'], 1)
    add(n + hours, np.roll(indexes['hydro_stock'], 1), -1)
    for name, sign in [('hydro_charge', -eta_h), ('hydro_release', 1), ('hydro_spill', 1)]:
        add(n + hours, indexes[name], sign)
    add(2 * n + hours, indexes['battery_stock'], 1)
    add(2 * n + hours, np.roll(indexes['battery_stock'], 1), -1)
    for name, sign in [('battery_charge', -eta_b), ('battery_discharge', 1 / eta_b)]:
        add(2 * n + hours, indexes[name], sign)
    matrix = coo_matrix((vv, (rr, cc)), shape=(3 * n, n * len(VARIABLES))).tocsr()
    rhs = np.concatenate([demand - generation, inflow, np.zeros(n)])
    maxima = {'hydro_release': power_h, 'hydro_charge': p['hydro_pump_gw'], 'hydro_stock': capacity_h, 'hydro_spill': np.inf, 'battery_charge': p['battery_gw'], 'battery_discharge': p['battery_gw'], 'battery_stock': p['battery_gwh'], 'imports': import_limit_gw, 'exports': p['export_limit_gw'], 'unserved': demand, 'curtailment': generation}
    upper = np.concatenate([np.broadcast_to(maxima[name], n) for name in VARIABLES])
    lower = np.zeros(len(upper))
    initial_fraction = p.get('initial_hydro_fill_fraction')
    if initial_fraction is not None:
        if not 0 <= initial_fraction <= 1:
            raise ValueError('Initial hydro fill fraction must lie between zero and one')
        lower[indexes['hydro_stock'][-1]] = capacity_h * initial_fraction
        upper[indexes['hydro_stock'][-1]] = capacity_h * initial_fraction
    objective = np.concatenate([np.full(n, config['objective_coefficients'].get(name, 0)) for name in VARIABLES])
    result = linprog(objective, A_eq=matrix, b_eq=rhs, bounds=np.column_stack([lower, upper]), method='highs', options={'dual_feasibility_tolerance': 1e-8, 'primal_feasibility_tolerance': 1e-8})
    if not result.success:
        raise RuntimeError(result.message)
    frame = pd.DataFrame({name: result.x[indexes[name]] for name in VARIABLES})
    error = matrix @ result.x - rhs
    tolerance = 1e-5
    assert np.max(np.abs(error)) < tolerance
    assert np.max(result.x - upper) < tolerance and np.min(result.x - lower) > -tolerance
    simultaneous_import_export = float(np.minimum(frame.imports, frame.exports).max())
    simultaneous_battery = float(np.minimum(frame.battery_charge, frame.battery_discharge).max())
    simultaneous_hydro = float(np.minimum(frame.hydro_charge, frame.hydro_release).max())
    assert max(simultaneous_import_export, simultaneous_battery, simultaneous_hydro) < tolerance
    audit = {'solver_success': True, 'solver_message': result.message, 'solve_seconds': perf_counter() - started, 'max_hourly_equation_residual_gwh': float(np.abs(error).max()), 'max_bound_violation': float(max(0, np.max(result.x - upper), np.max(lower - result.x))), 'hours': n, 'cyclic_hydro_initial_final_gwh': float(frame.hydro_stock.iloc[-1]), 'cyclic_battery_initial_final_gwh': float(frame.battery_stock.iloc[-1]), 'simultaneous_import_export_max_gw': simultaneous_import_export, 'simultaneous_battery_charge_discharge_max_gw': simultaneous_battery, 'simultaneous_hydro_charge_release_max_gw': simultaneous_hydro, 'reservoir_capacity_gwh': capacity_h, 'battery_capacity_gwh': p['battery_gwh'], 'initial_state_optimized': initial_fraction is None}
    return frame, audit


def summarize(frame, demand, supply):
    winter = local_index(frame.index).month.isin([10, 11, 12, 1, 2, 3])
    total = lambda column: float(frame[column].sum() / 1000)
    imports = total('imports')
    exports = total('exports')
    return {'hours': len(frame), 'final_demand_twh': float(demand.final_demand_gw.sum() / 1000), 'bus_demand_twh': float(demand.bus_demand_gw.sum() / 1000), 'peak_final_demand_gw': float(demand.final_demand_gw.max()), 'peak_bus_demand_gw': float(demand.bus_demand_gw.max()), 'winter_final_demand_twh': float(demand.final_demand_gw[winter].sum() / 1000), 'solar_available_twh': float(supply.solar_gw.sum() / 1000), 'nuclear_available_twh': float(supply.nuclear_gw.sum() / 1000), 'river_available_twh': float(supply.river_gw.sum() / 1000), 'reservoir_inflow_twh': float(supply.reservoir_inflow_gw.sum() / 1000), 'imports_twh': imports, 'exports_twh': exports, 'net_exports_twh': exports - imports, 'winter_imports_twh': float(frame.imports[winter].sum() / 1000), 'winter_exports_twh': float(frame.exports[winter].sum() / 1000), 'winter_net_imports_twh': float((frame.imports[winter] - frame.exports[winter]).sum() / 1000), 'unserved_twh': total('unserved'), 'shortage_hours': int((frame.unserved > 1e-5).sum()), 'peak_unserved_gw': float(frame.unserved.max()), 'winter_unserved_twh': float(frame.unserved[winter].sum() / 1000), 'hydro_charge_twh': total('hydro_charge'), 'hydro_release_twh': total('hydro_release'), 'hydro_spill_twh': total('hydro_spill'), 'battery_charge_twh': total('battery_charge'), 'battery_discharge_twh': total('battery_discharge'), 'curtailment_twh': total('curtailment'), 'hydro_storage_min_twh': float(frame.hydro_stock.min() / 1000), 'hydro_storage_max_twh': float(frame.hydro_stock.max() / 1000), 'adequate_in_deterministic_case': bool(frame.unserved.max() <= 1e-5), 'fuel_sovereignty_assessed': False}