#!/usr/bin/env python3

from copy import deepcopy
from math import log10
from warnings import filterwarnings
from gettext import install
from logging import basicConfig, getLogger
from pathlib import Path

from numpy import argmin, array, logspace
from rich.logging import RichHandler
from matplotlib.pyplot import figure

from backscattering_analyzer import Experiment, AcquisitionType, Sneb, Sweb, Sdb1, Sdb2  # pyright: ignore[reportMissingTypeStubs]
from backscattering_analyzer.bench import Bench  # pyright: ignore[reportMissingTypeStubs]
from backscattering_analyzer.utils import compute_light, get_factors  # pyright: ignore[reportMissingTypeStubs]
from science_signal.signal import Signal  # pyright: ignore[reportMissingTypeStubs]
from scipy.constants import pi  # pyright: ignore[reportMissingTypeStubs]

install(__name__)
filterwarnings("error")

basicConfig(
    level="INFO",
    format="%(message)s",
    datefmt="[%X]",
    handlers=[RichHandler()],
)
logger = getLogger(__name__)

sdb1: Bench = Sdb1()
sdb2: Bench = Sdb2()

excited: Experiment = Experiment(
    "2024_10_31", set([sdb1, sdb2]), Path("/home/demagny/data"), AcquisitionType.EXCITED
)

reference = Experiment(
    "2024_10_31",
    set([deepcopy(sdb1), deepcopy(sdb2)]),
    Path("/home/demagny/data"),
    AcquisitionType.REFERENCE,
)

excited.benches.remove(sdb1)
excited.projection.scatter_factor["SDB2"] = 1.62e-12
excited.projection.compute()
sensitivity_without_sdb2 = (
    excited.sensitivity.data.psd().sqrt() - excited.projection.data
)

"""
start of fit
"""

phase = 4 * pi / 1.064e-6
(
    factor_1,
    coupling_1,
    factor_2,
    coupling_2,
    signal_excited,
    signal_reference,
) = get_factors(
    sdb1,
    phase,
    sensitivity_without_sdb2,
    reference.sensitivity.data.psd().sqrt(),
    start=15,
    end=100,
)

scatter_min = 1e-20
scatter_max = 1e20

min: float = 0

for index in range(3):
    x = logspace(log10(scatter_max), log10(scatter_min), 1000)
    y = array(
        [
            sum(
                abs(
                    Signal(
                        factor_1.x,
                        compute_light(
                            x[i] / 1e-6,
                            factor_1,
                            coupling_1,
                            factor_2,
                            coupling_2,
                            phase,
                        ),
                    )
                    + signal_reference
                    - signal_excited
                ).y
            )
            for i in range(len(x))
        ]
    )
    if argmin(y) == 0:
        raise Exception(
            _("best scatter factor seems below current range: {}".format(scatter_min))
        )
    elif argmin(y) == len(x) - 1:
        raise Exception(
            _("best scatter factor seems above current range: {}".format(scatter_max))
        )
    else:
        scatter_min = x[argmin(y) - 1]
        scatter_max = x[argmin(y) + 1]

scatter_factor = (scatter_min + scatter_max) / 2

projection = (
    Signal(
        factor_1.x,
        compute_light(
            scatter_factor / 1e-6, factor_1, coupling_1, factor_2, coupling_2, phase
        ),
    )
    + reference.sensitivity.data.psd().sqrt()
)

"""
end of fit
"""

fig = figure(figsize=(20, 10))
fig.gca().loglog(
    sensitivity_without_sdb2.x,
    sensitivity_without_sdb2.y,
    label=_("excited sensitivity"),
)
fig.gca().loglog(
    projection.x, projection.y, label=_("projection + reference sensitivity")
)
fig.gca().legend()
fig.gca().grid(True, "both")
fig.gca().set_xlim(5, 100)
fig.gca().set_ylim(10e-28, 10e-18)
fig.gca().set_title(_("SDB1 (without fitting SDB2)"))
fig.gca().set_xlabel(_("frequencies (Hz)"))
fig.gca().set_ylabel(_("sensitivity ($\\frac{1} {\\sqrt {Hz}}$)"))
fig.savefig("SDB1.png")
with open("SDB1.csv", "w") as f:
    f.write(
        """# bench	value
SDB1	{}""".format(scatter_factor)
    )