import math
from dataclasses import dataclass
from typing import Optional


@dataclass
class RFStage:
    name: str
    gain_db: float
    noise_figure_db: float
    iip3_dbm: Optional[float] = None
    input_p1db_dbm: Optional[float] = None


lna = RFStage(name="LNA", gain_db=18.0, noise_figure_db=1.2, iip3_dbm=None, input_p1db_dbm=None)
rf_filter = RFStage(name="RF filter", gain_db=-2.0, noise_figure_db=2.0, iip3_dbm=None, input_p1db_dbm=None)
mixer = RFStage(name="Mixer", gain_db=-7.0, noise_figure_db=9.0, iip3_dbm=None, input_p1db_dbm=None)
if_amplifier = RFStage(name="IF amplifier", gain_db=24.0, noise_figure_db=3.0, iip3_dbm=None, input_p1db_dbm=None)

stages = [lna, rf_filter, mixer, if_amplifier]
input_power_dbm = -65.0


def db_to_linear(value_db):
    return 10 ** (value_db / 10)


def dbm_to_mw(value_dbm):
    return 10 ** (value_dbm / 10)


def mw_to_dbm(value_mw):
    return 10 * math.log10(value_mw)


def analyze_cascade(stages, input_power_dbm):
    total_gain_db = 0.0
    cumulative_gain_linear = 1.0
    total_noise_factor = 0.0
    inverse_iip3 = 0.0
    complete_iip3 = all(stage.iip3_dbm is not None for stage in stages)
    complete_p1db = all(stage.input_p1db_dbm is not None for stage in stages)
    p1db_candidates = []
    stage_results = []

    for index, stage in enumerate(stages):
        gain_before_db = total_gain_db
        stage_input_dbm = input_power_dbm + gain_before_db
        gain_linear = db_to_linear(stage.gain_db)
        noise_factor = db_to_linear(max(0.0, stage.noise_figure_db))

        if index == 0:
            total_noise_factor = noise_factor
        else:
            total_noise_factor += (noise_factor - 1) / cumulative_gain_linear

        if complete_iip3:
            inverse_iip3 += cumulative_gain_linear / dbm_to_mw(stage.iip3_dbm)

        if complete_p1db:
            p1db_candidates.append((stage.input_p1db_dbm - gain_before_db, stage.name))

        total_gain_db += stage.gain_db
        cumulative_gain_linear *= gain_linear
        stage_results.append({
            "name": stage.name,
            "input_dbm": stage_input_dbm,
            "output_dbm": input_power_dbm + total_gain_db,
            "cumulative_gain_db": total_gain_db,
        })

    cascade_iip3_dbm = mw_to_dbm(1 / inverse_iip3) if complete_iip3 else None
    cascade_oip3_dbm = cascade_iip3_dbm + total_gain_db if complete_iip3 else None
    limiting_p1db = min(p1db_candidates, key=lambda item: item[0]) if complete_p1db else None
    cascade_input_p1db_dbm = limiting_p1db[0] if limiting_p1db else None
    cascade_output_p1db_dbm = (
        cascade_input_p1db_dbm + total_gain_db - 1 if limiting_p1db else None
    )

    return {
        "total_gain_db": total_gain_db,
        "output_power_dbm": input_power_dbm + total_gain_db,
        "noise_figure_db": 10 * math.log10(total_noise_factor),
        "cascade_iip3_dbm": cascade_iip3_dbm,
        "cascade_oip3_dbm": cascade_oip3_dbm,
        "cascade_input_p1db_dbm": cascade_input_p1db_dbm,
        "cascade_output_p1db_dbm": cascade_output_p1db_dbm,
        "p1db_headroom_db": cascade_input_p1db_dbm - input_power_dbm if limiting_p1db else None,
        "limiting_stage": limiting_p1db[1] if limiting_p1db else None,
        "stages": stage_results,
    }


result = analyze_cascade(stages, input_power_dbm)

for stage in result["stages"]:
    print(f'{stage["name"]}: {stage["input_dbm"]:.2f} dBm in -> {stage["output_dbm"]:.2f} dBm out')

print(f'\nTotal gain: {result["total_gain_db"]:.2f} dB')
print(f'Output power: {result["output_power_dbm"]:.2f} dBm')
print(f'Noise figure: {result["noise_figure_db"]:.2f} dB')

if result["cascade_iip3_dbm"] is not None:
    print(f'Cascaded IIP3: {result["cascade_iip3_dbm"]:.2f} dBm')
    print(f'Cascaded OIP3: {result["cascade_oip3_dbm"]:.2f} dBm')
else:
    print('Cascaded IIP3/OIP3: enter IIP3 for every stage')

if result["cascade_input_p1db_dbm"] is not None:
    print(f'Input P1dB: {result["cascade_input_p1db_dbm"]:.2f} dBm')
    print(f'P1dB headroom: {result["p1db_headroom_db"]:.2f} dB')
    print(f'Limiting stage: {result["limiting_stage"]}')
else:
    print('Cascade P1dB: enter input P1dB for every stage')