#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# SPDX-License-Identifier: BSD-3-Clause
# test_gain_mode
# Copyright (C) 2025 Telecom Infra Project and GNPy contributors
# see AUTHORS.rst for a list of contributors

"""
checks behaviour of gain mode
- if all amps have their gains set, check that these gains are used, even if power_dbm or req_power change
- check that saturation is correct in gain mode

"""

from pathlib import Path
from numpy.testing import assert_array_equal, assert_allclose

import pytest
from gnpy.core.utils import automatic_nch, dbm2watt
from gnpy.core.network import build_network
from gnpy.tools.json_io import load_equipment, load_network, load_json
from gnpy.core.equipment import trx_mode_params
from gnpy.topology.request import PathRequest, compute_constrained_path, propagate


TEST_DIR = Path(__file__).parent
DATA_DIR = TEST_DIR / 'data'
EQPT_FILENAME = DATA_DIR / 'eqpt_config.json'
NETWORK_FILENAME = DATA_DIR / 'perdegreemeshTopologyExampleV2_auto_design_expected.json'
EXTRA_CONFIGS = {"std_medium_gain_advanced_config.json": load_json(DATA_DIR / "std_medium_gain_advanced_config.json")}


def pathrequest(pch_dbm, nb_channels):
    """create ref channel for defined power settings
    """
    params = {
        "power": dbm2watt(pch_dbm),
        "tx_power": dbm2watt(pch_dbm),
        "nb_channel": nb_channels,
        'request_id': None,
        'trx_type': None,
        'trx_mode': None,
        'source': None,
        'destination': None,
        'bidir': False,
        'nodes_list': [],
        'loose_list': [],
        'format': '',
        'baud_rate': None,
        'bit_rate': None,
        'roll_off': None,
        'OSNR': None,
        'penalties': None,
        'path_bandwidth': None,
        'effective_freq_slot': None,
        'f_min': None,
        'f_max': None,
        'spacing': None,
        'min_spacing': None,
        'cost': None,
        'equalization_offset_db': None,
        'tx_osnr': None
    }
    return PathRequest(**params)


def net_setup(equipment):
    """Common setup for tests: builds network, equipment
    """
    network = load_network(NETWORK_FILENAME, equipment)
    spectrum = equipment['SI']['default']
    p_db = spectrum.power_dbm
    nb_channels = automatic_nch(spectrum.f_min, spectrum.f_max, spectrum.spacing)
    build_network(network, equipment, pathrequest(p_db, nb_channels))
    return network


def create_rq(equipment, srce, dest, bdir, nd_list, ls_list, mode, power_dbm):
    """Create the usual request list according to parameters
    """
    params = {
        'request_id': 'test_request',
        'source': srce,
        'bidir': bdir,
        'destination': dest,
        'trx_type': 'Voyager',
        'trx_mode': mode,
        'format': mode,
        'nodes_list': nd_list,
        'loose_list': ls_list,
        'effective_freq_slot': None,
        'path_bandwidth': 100000000000.0,
        'spacing': 50e9 if mode == 'mode 1' else 75e9,
        'power': dbm2watt(power_dbm),
        'tx_power': dbm2watt(power_dbm)
    }
    trx_params = trx_mode_params(equipment, params['trx_type'], params['trx_mode'], True)
    params.update(trx_params)
    f_min = params['f_min']
    f_max_from_si = params['f_max']
    params['nb_channel'] = automatic_nch(f_min, f_max_from_si, params['spacing'])
    return PathRequest(**params)


@pytest.mark.parametrize("power_dbm", [0, -2, 3])
@pytest.mark.parametrize("req_power", [1e-3, 0.5e-3, 2e-3])
def test_gain_mode(req_power, power_dbm):
    """ Gains are all set on the selected path, so that since the design is made for 0dBm,
    in gain mode, whatever the value of equipment power_dbm or request power, the network is unchanged
    and the propagation remains the same as for power mode and 0dBm
    """
    equipment = load_equipment(EQPT_FILENAME, EXTRA_CONFIGS)
    network = net_setup(equipment)
    req = create_rq(equipment, 'trx Brest_KLA', 'trx Rennes_STA', False,
                    ['Edfa0_roadm Brest_KLA', 'roadm Lannion_CAS', 'trx Rennes_STA'],
                    ['STRICT', 'STRICT', 'STRICT'], 'mode 1', 0)
    path = compute_constrained_path(network, req)
    # Propagation in power_mode
    infos_expected = propagate(path, req, equipment)
    # Now set to gain mode
    setattr(equipment['Span']['default'], 'power_mode', False)
    setattr(equipment['SI']['default'], 'power_dbm', power_dbm)
    req.power = req_power
    network2 = net_setup(equipment)
    path2 = compute_constrained_path(network2, req)
    infos_actual = propagate(path2, req, equipment)

    assert_array_equal(infos_expected.baud_rate, infos_actual.baud_rate)
    assert_allclose(infos_expected.signal, infos_actual.signal, rtol=1e-14)
    assert_allclose(infos_expected.nli, infos_actual.nli, rtol=1e-14)
    assert_allclose(infos_expected.ase, infos_actual.ase, rtol=1e-14)
    assert_array_equal(infos_expected.roll_off, infos_actual.roll_off)
    assert_array_equal(infos_expected.chromatic_dispersion, infos_actual.chromatic_dispersion)
    assert_array_equal(infos_expected.pmd, infos_actual.pmd)
    assert_array_equal(infos_expected.channel_number, infos_actual.channel_number)
    assert_array_equal(infos_expected.number_of_channels, infos_actual.number_of_channels)