from collections.abc import Sequence
from pathlib import Path
from typing import Any
from typing import Optional
from typing import Tuple
from typing import Union

import numpy
import numpy as np
from scipy.optimize import curve_fit
from scipy.stats import linregress

from .hdf5 import get_measurement_data
from .hdf5 import get_scan_motor
from .hdf5 import get_scan_start_time
from .hdf5 import get_scan_title
from .utils import distance_from_COR
from .utils import logger


def gauss(
    x: numpy.ndarray, H: float, A: float, x0: float, sigma: float
) -> numpy.ndarray:
    return H + A * numpy.exp(-((x - x0) ** 2) / (2 * sigma**2))


def gauss_fit(x: numpy.ndarray, y: numpy.ndarray) -> numpy.ndarray:
    sum_y = sum(y)
    if sum_y == 0:
        sum_y = 1
    mean = sum(x * y) / sum_y
    sigma = numpy.sqrt(sum(y * (x - mean) ** 2) / sum_y)
    popt, pcov = curve_fit(gauss, x, y, p0=[min(y), max(y), mean, sigma])
    return popt


def linear_fit(x: numpy.ndarray, y: numpy.ndarray) -> Any:
    return linregress(x, y)


def CDI_align_estimate_COR(
    filename: Union[str, Path],
    scan_nos: Sequence[int],
    axis_name: str,
    rc_axis_name: str,
    counter: str,
    retry_timeout: Optional[int] = None,
    retry_period: Optional[int] = None,
) -> Tuple[list, list, list, list, numpy.ndarray, numpy.ndarray, numpy.ndarray]:
    """
    From a number of alignment scans at different points on a rocking curve
    determine the location of the COR relative to the current position.
    """
    fitpars = np.recarray(
        (len(scan_nos),),
        dtype=[
            ("scan_no", "<f8"),
            ("start_time", "<f8"),
            ("rc_angle", "<f8"),
            ("label", "|S50"),
            ("x0", "<f8"),
            ("sigma", "<f8"),
            ("H", "<f8"),
            ("A", "<f8"),
            ("FWHM", "<f8"),
        ],
    )
    xs = []
    intensities = []
    intensity_fits = []
    labels = []

    for i, scan_no in enumerate(scan_nos):
        logger.info(f"scan_no: {scan_no}")

        intensity = get_measurement_data(
            scan_no,
            filename,
            counter,
            retry_timeout=retry_timeout,
            retry_period=retry_period,
        )
        x = get_measurement_data(
            scan_no,
            filename,
            axis_name,
            retry_timeout=retry_timeout,
            retry_period=retry_period,
        )
        label = get_scan_title(
            scan_no,
            filename,
            retry_period=retry_period,
        )

        H, A, x0, sigma = gauss_fit(x, intensity)
        FWHM = 2.35482 * sigma

        # get the motor position for the RC scan
        rc_angle = get_scan_motor(
            scan_no,
            filename,
            rc_axis_name,
            retry_timeout=retry_timeout,
            retry_period=retry_period,
        )
        start_time = get_scan_start_time(
            scan_no,
            filename,
            retry_timeout=retry_timeout,
            retry_period=retry_period,
        )
        fitpars[i] = (scan_no, start_time, rc_angle, label, x0, sigma, H, A, FWHM)
        xs.append(x)
        intensities.append(intensity)
        intensity_fit = gauss(x, *gauss_fit(x, intensity))
        intensity_fits.append(intensity_fit)
        labels.append(label)

    offset_x, offset_y = fitpars["rc_angle"], fitpars["x0"]
    res = linear_fit(offset_x, offset_y)
    logger.info(f"R-squared: {res.rvalue**2:.6f}")
    offset_fit = res.intercept + offset_x * res.slope

    logger.info(
        f"The observed offset per degree {rc_axis_name} for {axis_name} : {res.slope}"
    )
    logger.info("apply this to your scans with d2scans or")
    logger.info(".... you could move the crystal closer to the rotation axis with")
    distance_from_COR(1, res.slope)
    # calculate how far away the center of rotation is.

    return xs, intensities, intensity_fits, labels, offset_x, offset_y, offset_fit
