from time import perf_counter

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


def solve(demand, generation, inflow, options, road=None, heat=None, index=None):
    started = perf_counter()
    demand, generation, inflow = [np.asarray(a, dtype=float) for a in (demand, generation, inflow)]
    n = len(demand)
    if not n or any(len(a) != n or not np.isfinite(a).all() or a.min() < 0 for a in (demand, generation, inflow)):
        raise ValueError('Invalid hourly input')
    stores = options['pumped_stores']
    names = ['natural_release', 'natural_stock', 'natural_spill', 'battery_charge', 'battery_discharge', 'battery_stock', 'imports', 'exports', 'unserved', 'curtailment']
    for store in stores:
        names += [store['id'] + '_' + suffix for suffix in ('charge', 'discharge', 'stock')]
    flex_fraction = options.get('road_flexible_fraction', 0)
    if flex_fraction:
        if index is None or road is None:
            raise ValueError('Road flexibility requires dated road load')
        names.append('road_flex')
    thermal_power = options.get('thermal_shift_gw', 0)
    if thermal_power:
        if heat is None:
            raise ValueError('Thermal shifting requires heating load')
        names += ['thermal_charge', 'thermal_discharge', 'thermal_stock']
    ix = {name: i * n + np.arange(n) for i, name in enumerate(names)}
    p_ix, e_ix = n * len(names), n * len(names) + 1
    size = e_ix + 1
    er, ec, ev, rhs = [], [], [], list(demand - generation)
    ur, uc, uv, ub = [], [], [], []
    hours = np.arange(n)

    def eq(rows, cols, values):
        er.extend(np.atleast_1d(rows))
        ec.extend(np.atleast_1d(cols))
        ev.extend(np.broadcast_to(values, len(np.atleast_1d(rows))))

    def inequal(rows, cols, values):
        ur.extend(np.atleast_1d(rows))
        uc.extend(np.atleast_1d(cols))
        uv.extend(np.broadcast_to(values, len(np.atleast_1d(rows))))

    for name, sign in [('natural_release', 1), ('battery_charge', -1), ('battery_discharge', 1), ('imports', 1), ('exports', -1), ('unserved', 1), ('curtailment', -1)]:
        eq(hours, ix[name], sign)
    for store in stores:
        eq(hours, ix[store['id'] + '_charge'], -1)
        eq(hours, ix[store['id'] + '_discharge'], 1)
    row_start = n
    eq(row_start + hours, ix['natural_stock'], 1)
    eq(row_start + hours, np.roll(ix['natural_stock'], 1), -1)
    eq(row_start + hours, ix['natural_release'], 1)
    eq(row_start + hours, ix['natural_spill'], 1)
    rhs.extend(inflow)
    row_start += n
    for label, eta in [('battery', options['battery_rte'])] + [(s['id'], s['roundtrip_efficiency']) for s in stores]:
        eq(row_start + hours, ix[label + '_stock'], 1)
        eq(row_start + hours, np.roll(ix[label + '_stock'], 1), -1)
        eq(row_start + hours, ix[label + '_charge'], -eta)
        eq(row_start + hours, ix[label + '_discharge'], 1)
        rhs.extend(np.zeros(n))
        row_start += n
    road_upper = None
    if flex_fraction:
        road = np.asarray(road, dtype=float)
        if len(road) != n or not 0 <= flex_fraction <= 1 or np.any(road > demand):
            raise ValueError('Invalid road flexibility')
        rhs[:n] = np.asarray(rhs[:n]) - flex_fraction * road
        eq(hours, ix['road_flex'], -1)
        local = index.tz_convert('Europe/Zurich')
        dates = np.asarray(local.date)
        road_upper = np.zeros(n)
        available = (local.hour <= 6) | ((local.hour >= 10) & (local.hour <= 16)) | (local.hour >= 20)
        for date in sorted(set(dates)):
            mask = dates == date
            positions = np.flatnonzero(mask)
            energy = float(road[mask].sum() * flex_fraction)
            road_upper[mask] = energy / mask.sum() * 3 * available[mask]
            if road_upper[mask].sum() + 1e-8 < energy:
                raise ValueError('Daily flexible charging capacity insufficient')
            eq(np.full(len(positions), row_start), ix['road_flex'][positions], 1)
            rhs.append(energy)
            row_start += 1
    if thermal_power:
        heat = np.asarray(heat, dtype=float)
        eq(hours, ix['thermal_charge'], -1)
        eq(hours, ix['thermal_discharge'], 1)
        eq(row_start + hours, ix['thermal_stock'], 1)
        eq(row_start + hours, np.roll(ix['thermal_stock'], 1), -1)
        eq(row_start + hours, ix['thermal_charge'], -.95)
        eq(row_start + hours, ix['thermal_discharge'], 1)
        rhs.extend(np.zeros(n))
        row_start += n
    for label in ['battery_charge', 'battery_discharge', 'battery_stock']:
        start = len(ub)
        inequal(start + hours, ix[label], 1)
        inequal(start + hours, np.full(n, e_ix if label == 'battery_stock' else p_ix), -1)
        ub.extend(np.zeros(n))
    duration = options.get('battery_duration_hours')
    if duration is not None:
        eq([row_start, row_start], [e_ix, p_ix], [1, -duration])
        rhs.append(0)
        row_start += 1
    bounds = np.zeros((size, 2))
    bounds[:, 1] = np.inf
    caps = {'natural_release': options['natural_turbine_gw'], 'natural_stock': options['natural_capacity_gwh'], 'imports': options.get('import_gw', 0), 'exports': options.get('export_gw', 8), 'unserved': demand}
    for name, cap in caps.items():
        bounds[ix[name], 1] = cap
    if road_upper is not None:
        bounds[ix['road_flex'], 1] = road_upper
    if thermal_power:
        bounds[ix['thermal_charge'], 1] = thermal_power
        bounds[ix['thermal_discharge'], 1] = np.minimum(thermal_power, heat)
        bounds[ix['thermal_stock'], 1] = thermal_power * options.get('thermal_shift_hours', 6)
    for store in stores:
        for suffix, cap in [('charge', store['pump_gw']), ('discharge', store['turbine_gw']), ('stock', store['deliverable_energy_gwh'])]:
            bounds[ix[store['id'] + '_' + suffix], 1] = cap
    bounds[p_ix, 1] = options.get('battery_max_gw', 20)
    bounds[e_ix, 1] = options.get('battery_max_gwh', 5000)
    for name, position in [('battery_fixed_gw', p_ix), ('battery_fixed_gwh', e_ix)]:
        if name in options:
            bounds[position] = options[name]
    if 'natural_initial_fraction' in options:
        bounds[ix['natural_stock'][-1]] = options['natural_initial_fraction'] * options['natural_capacity_gwh']
    objective = np.zeros(size)
    objective[ix['unserved']] = 1e6
    objective[ix['imports']] = .001
    objective[ix['natural_spill']] = 1e-10
    for name in names:
        if name.endswith('_charge') or name == 'natural_release':
            objective[ix[name]] = 1e-9
    objective[p_ix] = options.get('battery_power_chf_per_kw', 600) / 1000
    objective[e_ix] = options.get('battery_energy_chf_per_kwh', 225) / 1000
    ae = coo_matrix((ev, (er, ec)), shape=(len(rhs), size)).tocsr()
    au = coo_matrix((uv, (ur, uc)), shape=(len(ub), size)).tocsr()
    result = linprog(objective, A_eq=ae, b_eq=rhs, A_ub=au, b_ub=ub, bounds=bounds, method='highs')
    if not result.success:
        raise RuntimeError(result.message)
    capacity_optimum = result.x[[p_ix, e_ix]].copy()
    bounds[p_ix] = capacity_optimum[0]
    bounds[e_ix] = capacity_optimum[1]
    dispatch_objective = np.zeros(size)
    dispatch_objective[ix['unserved']] = 1e6
    dispatch_objective[ix['imports']] = 100
    dispatch_objective[ix['natural_release']] = .0001
    dispatch_objective[ix['natural_spill']] = .000001
    dispatch_objective[ix['curtailment']] = .000001
    for name in names:
        if name.endswith('_charge'):
            dispatch_objective[ix[name]] = .02
        if name.endswith('_discharge'):
            dispatch_objective[ix[name]] = .0001
    result = linprog(dispatch_objective, A_eq=ae, b_eq=rhs, A_ub=au, b_ub=ub, bounds=bounds, method='highs')
    if not result.success:
        raise RuntimeError(result.message)
    x = result.x
    residual = float(np.max(np.abs(ae @ x - rhs)))
    bound_violation = float(max(0, np.max(x - bounds[:, 1]), np.max(bounds[:, 0] - x)))
    inequality = float(max(0, np.max(au @ x - ub)))
    if max(residual, bound_violation, inequality) > 2e-5:
        raise RuntimeError('Conservation or capacity validation failed')
    frame = pd.DataFrame({name: x[ix[name]] for name in names}, index=index)
    annual = {name + '_gwh': float(frame[name].sum()) for name in names if not name.endswith('_stock')}
    annual.update(battery_gw=float(x[p_ix]), battery_deliverable_gwh=float(x[e_ix]), battery_initial_capex_chf_billion=float(objective[p_ix] * x[p_ix] + objective[e_ix] * x[e_ix]), unserved_peak_gw=float(frame.unserved.max()), shortage_hours=int((frame.unserved > 1e-5).sum()), adequate=bool(frame.unserved.max() <= 1e-5), demand_gwh=float(demand.sum()), generation_gwh=float(generation.sum()), natural_inflow_gwh=float(inflow.sum()), solve_seconds=perf_counter() - started)
    annual['audit'] = dict(max_equation_residual_gwh=residual, max_bound_violation=bound_violation, max_inequality_violation=inequality, hours=n, initial_natural_gwh=float(frame.natural_stock.iloc[-1]), initial_battery_gwh=float(frame.battery_stock.iloc[-1]))
    annual['battery_annual_output_equivalent_cycles'] = annual['battery_discharge_gwh'] / x[e_ix] if x[e_ix] > 1e-7 else 0
    return frame, annual