numba in economics

from typing import TYPE_CHECKING, NamedTuple

import numpy as np

from numba import njit

from numpy.typing import NDArray

from pythermogis.workflow.utc.doublet_utils import get_along_hole_length

    total_opex_ts = np.zeros(n_years)

    total_capex_ts = np.zeros(n_years)

    shifted_heat_power = heat_power_produced.copy()

    shifted_heat_power = np.array(heat_power_produced, dtype=np.float64)

    shift_time_series(shifted_heat_power, props.drilling_time)

    capex_well = calculate_capex_for_wells(props, ah_length_array, stimulation_capex)

    allocate_capex_for_wells(props, total_capex_ts, capex_well)

    capex_well = calculate_capex_for_wells(

        inj_depth=ah_length_array[0],

        prod_depth=ah_length_array[0],

        stimulation_capex=stimulation_capex,

        well_cost_z2=props.well_cost_z2,

        well_cost_z=props.well_cost_z,

        well_cost_const=props.well_cost_const,

        well_cost_scaling=props.well_cost_scaling,

    )

    allocate_capex_for_wells(props.drilling_time, total_capex_ts, capex_well)

    hp_added_power = get_max_hp_added_power(

        initial_hp_added_power, hp_added_power_years

    hp_years_arr = (

        np.array(hp_added_power_years, dtype=np.float64)

        if hp_added_power_years is not None

        else np.zeros(0, dtype=np.float64)

    )

    hp_added_power = get_max_hp_added_power(initial_hp_added_power, hp_years_arr)

    hp_after_hp = calculate_hp_added_power_after_heat_pump(

        props, hp_added_power, hp_cop

        hp_added_power, hp_cop, props.hp_alternative_heating_price

    )

    installation_mw = calculate_installation_mw(

        float(shifted_heat_power[-1]), hp_added_power

    )

    installation_mw = calculate_installation_mw(shifted_heat_power, hp_added_power)

    total_capex_1year = calculate_total_capex_1year(

        props, total_capex_ts, hp_after_hp, installation_mw

        last_year=max(props.drilling_time - 1, 0),

        total_capex_ts=total_capex_ts,

        capex_const=props.capex_const,

        hp_capex=props.hp_capex,

        capex_variable=props.capex_variable,

        capex_contingency=props.capex_contingency,

        hp_after_hp=hp_after_hp,

        installation_mw=installation_mw,

    )

    sum_capex = process_annual_data(

            series[i] = 0

@njit

def calculate_capex_for_wells(

    props: UTCConfiguration, ah_length_array: list[float], stimulation_capex: float

    inj_depth: float,

    prod_depth: float,

    stimulation_capex: float,

    well_cost_z2: float,

    well_cost_z: float,

    well_cost_const: float,

    well_cost_scaling: float,

) -> float:

    inj_depth = ah_length_array[0]

    prod_depth = ah_length_array[0]

    capex = (

        calculate_well_cost(props, inj_depth)

        + calculate_well_cost(props, prod_depth)

        calculate_well_cost(well_cost_z2, well_cost_z, well_cost_const, inj_depth)

        + calculate_well_cost(well_cost_z2, well_cost_z, well_cost_const, prod_depth)

        + stimulation_capex

    )

    return capex * props.well_cost_scaling

    return capex * well_cost_scaling

def calculate_well_cost(props: UTCConfiguration, depth: float) -> float:

    return (

        props.well_cost_z2 * depth**2 + props.well_cost_z * depth

    ) * 1e-6 + props.well_cost_const

@njit

def calculate_well_cost(

    well_cost_z2: float, well_cost_z: float, well_cost_const: float, depth: float

) -> float:

    return (well_cost_z2 * depth**2 + well_cost_z * depth) * 1e-6 + well_cost_const

@njit

def allocate_capex_for_wells(

    props: UTCConfiguration, total_capex_ts: NDArray[np.float64], capex_well: float

):

    drilling_years = props.drilling_time

    drilling_years: int, total_capex_ts: np.ndarray, capex_well: float

) -> None:

    yearly_cost = capex_well / max(drilling_years, 1)

    for y in range(drilling_years):

        total_capex_ts[y] += yearly_cost

def get_max_hp_added_power(initial: float, hp_power_years: list[float] | None) -> float:

    if hp_power_years is not None:

        return max(hp_power_years)

@njit

def get_max_hp_added_power(initial: float, hp_power_years: np.ndarray) -> float:

    if hp_power_years.size == 0:

        return initial

    max_val = hp_power_years[0]

    for i in range(hp_power_years.size):

        if hp_power_years[i] > max_val:

            max_val = hp_power_years[i]

    return max_val

@njit

def calculate_hp_added_power_after_heat_pump(

    props: UTCConfiguration, hp_added_power: float, hp_cop: float

    hp_added_power: float, hp_cop: float, hp_alternative_heating_price: float

) -> float:

    alt_price = props.hp_alternative_heating_price

    return hp_added_power if alt_price < 0 else hp_added_power * (1 + 1 / (hp_cop - 1))

    if hp_alternative_heating_price < 0:

        return hp_added_power

    return hp_added_power * (1 + 1.0 / (hp_cop - 1.0))

@njit

def calculate_installation_mw(

    heat_power_produced: list[float], hp_added_power: float

    last_heat_power_value: float, hp_added_power: float

) -> float:

    last_val = heat_power_produced[-1]

    return max(last_val - hp_added_power, 0)

    diff = last_heat_power_value - hp_added_power

    return diff if diff > 0 else 0.0

@njit

def calculate_total_capex_1year(

    props: UTCConfiguration,

    total_capex_ts: NDArray[np.float64],

    last_year: int,

    total_capex_ts: np.ndarray,

    capex_const: float,

    hp_capex: float,

    capex_variable: float,

    capex_contingency: float,

    hp_after_hp: float,

    installation_mw: float,

) -> float:

    last_year = max(props.drilling_time - 1, 0)

    total_capex_ts[last_year] += (

        props.capex_const

        + (props.hp_capex * hp_after_hp + props.capex_variable * installation_mw) * 1e-3

        capex_const + (hp_capex * hp_after_hp + capex_variable * installation_mw) * 1e-3

    )

    total_capex_ts[last_year] *= 1 + props.capex_contingency / 100

    total_capex_ts[last_year] *= 1.0 + capex_contingency / 100.0

    return total_capex_ts[last_year]

        # Variable OPEX

        variable_opex[year] = calculate_variable_opex(

            props,

            season_factor,

            props.elec_purchase_price,

            pump_power_required,

            heat_power_per_year[year],

            year_hp_cop,

            hp_after_hp,

            props.opex_per_energy,

        )

        # Fixed OPEX

        fixed_opex[year] = calculate_fixed_opex(

            props,

            props.opex_base,

            props.hp_opex,

            props.opex_per_power,

            hp_after_hp,

            installation_mw,

            total_capex_1year,

            props.opex_per_capex,

        )

        total_opex_ts[year] = variable_opex[year] + fixed_opex[year]

    return sum_capex

@njit

def calculate_variable_opex(

    props: UTCConfiguration,

    season_factor: float,

    elec_price: float,

    pump_power: float,

    heat_power_gj_per_year: float,

    hp_cop: float,

    hp_after_hp: float,

    opex_per_energy: float,

) -> float:

    pump_cost = (

        -(pump_power * 1e3)

        * 1e-6

    )

    opex_per_energy = (

        -(1 / 0.36)

        * props.opex_per_energy

        * 1e-2

        * heat_power_gj_per_year

        * 1e4

        / 36

        * 1e-6

    opex_energy_cost = (

        -(1 / 0.36) * opex_per_energy * 1e-2 * heat_power_gj_per_year * 1e4 / 36 * 1e-6

    )

    heat_pump_cost = 0.0

            * 1e-6

        )

    return pump_cost + parasitic_cost + heat_pump_cost + opex_per_energy

    return pump_cost + parasitic_cost + heat_pump_cost + opex_energy_cost

@njit

def calculate_fixed_opex(

    props: UTCConfiguration,

    opex_base: float,

    hp_opex: float,

    opex_per_power: float,

    hp_after_hp: float,

    installation_mw: float,

    total_capex: float,

    opex_per_capex: float,

) -> float:

    return (

        -props.opex_base / 1e6

        - (props.hp_opex * hp_after_hp + props.opex_per_power * installation_mw) * 1e-3

        - total_capex * props.opex_per_capex / 100

        -opex_base / 1e6

        - (hp_opex * hp_after_hp + opex_per_power * installation_mw) * 1e-3

        - total_capex * opex_per_capex / 100.0

    )