import math

import numpy as np
from numba import njit


PARAMETERS = {
    'coilInductanceHenries': 2.,
    'coilResistanceOhms': 10000.,
    'coilCapacitanceFarads': 50e-12,
    'cableCapacitanceFarads': 400e-12,
    'coreLossResistanceOhms': 1e6,
    'volumeLoadResistanceOhms': 500000.,
    'amplifierLoadResistanceOhms': 1e6,
}


def components(parameters=None):
    values = PARAMETERS if parameters is None else parameters
    inductance = float(values['coilInductanceHenries'])
    resistance = float(values['coilResistanceOhms'])
    capacitance = float(values['coilCapacitanceFarads']) + float(values['cableCapacitanceFarads'])
    loads = [float(values[name]) for name in
             ('coreLossResistanceOhms', 'volumeLoadResistanceOhms', 'amplifierLoadResistanceOhms')]
    if not math.isfinite(inductance) or inductance <= 0 or not math.isfinite(resistance) or resistance < 0:
        raise ValueError('Invalid pickup coil')
    if not math.isfinite(capacitance) or capacitance <= 0 or any(value <= 0 or math.isnan(value) for value in loads):
        raise ValueError('Invalid pickup capacitance/load')
    conductance = sum(1 / value for value in loads)
    return inductance, resistance, capacitance, conductance


def response(frequencies, parameters=None, sample_rate=None):
    inductance, resistance, capacitance, conductance = components(parameters)
    frequencies = np.asarray(frequencies, dtype=np.float64)
    angular = 2 * np.pi * frequencies if sample_rate is None else 2 * sample_rate * np.tan(np.pi * frequencies / sample_rate)
    return 1 / (1 + resistance * conductance - angular ** 2 * inductance * capacitance
                + 1j * angular * (inductance * conductance + resistance * capacitance))


@njit(cache=True)
def add_compensated(total, compensation, addition):
    corrected = addition - compensation
    result = total + corrected
    return result, (result - total) - corrected


@njit(cache=True)
def integrate(source, rate, inductance, resistance, capacitance, conductance,
              initial_current, initial_voltage, stride):
    time_step = 1 / rate
    angular = 1 / math.sqrt(inductance * capacitance)
    current_loss = resistance / inductance
    voltage_loss = conductance / capacitance
    current_diagonal = 1 + time_step * current_loss / 2
    voltage_diagonal = 1 + time_step * voltage_loss / 2
    cross = time_step * angular / 2
    determinant = current_diagonal * voltage_diagonal + cross * cross
    current_state = math.sqrt(inductance) * initial_current
    voltage_state = math.sqrt(capacitance) * initial_voltage
    initial_energy = .5 * (current_state ** 2 + voltage_state ** 2)
    output = np.empty(len(source))
    ledger = np.empty((2 + len(source) // stride, 7))
    ledger[0] = np.array([0., initial_energy, 0., 0., 0., 0., 0.])
    samples = 1
    supplied = 0.
    coil_loss = 0.
    load_loss = 0.
    compensation = np.zeros(3)
    previous_source = 0.
    maximum_residual = 0.
    maximum_current = abs(initial_current)
    for frame in range(len(source)):
        midpoint_source = .5 * (source[frame] + previous_source)
        current_rhs = (1 - time_step * current_loss / 2) * current_state - cross * voltage_state
        current_rhs += time_step * midpoint_source / math.sqrt(inductance)
        voltage_rhs = cross * current_state + (1 - time_step * voltage_loss / 2) * voltage_state
        next_current = (voltage_diagonal * current_rhs - cross * voltage_rhs) / determinant
        next_voltage = (cross * current_rhs + current_diagonal * voltage_rhs) / determinant
        midpoint_current = .5 * (current_state + next_current) / math.sqrt(inductance)
        midpoint_voltage = .5 * (voltage_state + next_voltage) / math.sqrt(capacitance)
        supplied, compensation[0] = add_compensated(supplied, compensation[0], time_step * midpoint_source * midpoint_current)
        coil_loss, compensation[1] = add_compensated(coil_loss, compensation[1], time_step * resistance * midpoint_current ** 2)
        load_loss, compensation[2] = add_compensated(load_loss, compensation[2], time_step * conductance * midpoint_voltage ** 2)
        current_state = next_current
        voltage_state = next_voltage
        previous_source = source[frame]
        output[frame] = voltage_state / math.sqrt(capacitance)
        maximum_current = max(maximum_current, abs(current_state / math.sqrt(inductance)))
        if frame % stride == stride - 1 or frame == len(source) - 1:
            stored = .5 * (current_state ** 2 + voltage_state ** 2)
            residual = stored - initial_energy - supplied + coil_loss + load_loss
            maximum_residual = max(maximum_residual, abs(residual))
            ledger[samples] = np.array([(frame + 1) / rate, stored, supplied, coil_loss,
                                        load_loss, residual, current_state / math.sqrt(inductance)])
            samples += 1
    return output, ledger[:samples], maximum_residual, maximum_current


def simulate(source, sample_rate, parameters=None, *, initial_current=0., initial_voltage=0., ledger_rate=100):
    source = np.asarray(source, dtype=np.float64)
    if source.ndim != 1 or not len(source) or not np.isfinite(source).all():
        raise ValueError('Invalid pickup EMF source')
    if not isinstance(sample_rate, int) or sample_rate < 1000 or ledger_rate <= 0:
        raise ValueError('Invalid circuit sampling controls')
    if not math.isfinite(initial_current) or not math.isfinite(initial_voltage):
        raise ValueError('Invalid initial electrical state')
    values = components(parameters)
    output, ledger, residual, maximum_current = integrate(source, sample_rate, *values,
                                                         initial_current, initial_voltage,
                                                         max(1, sample_rate // ledger_rate))
    if not np.isfinite(output).all() or not np.isfinite(ledger).all():
        raise FloatingPointError('Non-finite pickup circuit')
    return {'outputVolts': output, 'ledger': ledger,
            'maximumEnergyResidualJoules': float(residual),
            'maximumCoilCurrentAmperes': float(maximum_current), 'sampleRate': sample_rate}