"""
Experiment definitions for SSX collections:
    - extruder
    - fixed target
    - 3D grid scan
"""
from __future__ import annotations

import logging
from pathlib import Path
from typing import Dict, List, Tuple

import numpy as np

from ..nxs_utils import Axis, TransformationType
from ..nxs_utils.ScanUtils import calculate_scan_points
from .beamline_utils import PumpProbe
from .SSX_chip import (
    Chip,
    compute_goniometer,
    fullchip_blocks_conversion,
    read_chip_map,
)

__all__ = ["run_extruder", "run_fixed_target", "run_3D_grid_scan"]

# Define a logger object
logger = logging.getLogger("nexgen.SSX.run_expt")


def run_extruder(
    goniometer_axes: List[Axis],
    num_imgs: int,
    pump_probe: PumpProbe,
    osc_axis: str = "omega",
) -> Tuple[List, Dict, Dict]:
    """Run the goniometer computations for an extruder experiment.

    Args:
        goniometer_axes (List[Axis]): List of goniometer axes for current beamline.
        num_imgs (int): Total number of images.
        pump_probe (PumpProbe): Pump probe parameters.
        osc_axis: Defines which axis is considered the "moving" one. Defaults to omega.

    Returns:
        Tuple[List, Dict, Dict]:
            goniometer_axes: updated goniometer_axes list with actual values from the scan.
            SCAN: dictionary with oscillation scan axis values.
            pump_info: updated pump probe information.
    """
    logger.debug("Running an extruder experiment.")

    logger.debug("All axes are fixed, setting increments to 0.0 and starts == ends.")
    for ax in goniometer_axes:
        # Sanity check that no increment is greater than 0.0
        if ax.transformation_type == "rotation":
            ax.increment = 0.0

    # Identify the "oscillation axis"
    osc_idx = [n for n, ax in enumerate(goniometer_axes) if ax.name == osc_axis][0]
    goniometer_axes[osc_idx].num_steps = num_imgs

    # Calculate scan
    logger.debug(
        "Getting 'oscillation scan': roation axis not moving, same value for each image as."
    )
    SCAN = calculate_scan_points(
        goniometer_axes[osc_idx], rotation=True, tot_num_imgs=num_imgs
    )

    pump_info = pump_probe.to_dict()
    logger.debug("Removing pump_repeat from pump probe necessary information.")
    pump_info.pop("pump_repeat")

    return goniometer_axes, SCAN, pump_info


def run_fixed_target(
    goniometer_axes: List[Axis],
    chip_info: Dict[str, List],
    chipmap: Path | str,
    pump_probe: PumpProbe,
    scan_axes: List[str, str] = ["sam_y", "sam_x"],
) -> Tuple[Dict, Dict]:
    """Run the goniometer computations for a fixed-target experiment.

    Args:
        goniometer_axes (List[Axis]): List of goniometer axes for current beamline.
        chip_info (Dict[str, List]): General information about the chip: number and size of blocks, \
            size and step of each window, start positions, number of exposures.
        chipmap (Path | str): Path to .map file. If None is passed, assumes a fullchip.
        pump_probe (PumpProbe): Pump probe parameters.
        scan_axes (List[str, str], optional): List of scan axes, in order slow,fast. \
            Defaults to ["sam_y", "sam_x"].

    Raises:
        ValueError: If one or both of the axes names passed as input are not part of the goniometer axes.
        ValueError:if chip_info hasn't been passed or is an empty dictionary.

    Returns:
        Tuple[Dict, Dict]:
            SCAN: Dictionary with grid scan values.
            pump_info: Updated pump probe information.
    """
    logger.info("Running a fixed target experiment.")

    # Check that the axes for the scan make sense
    check_list = [n for n, ax in enumerate(goniometer_axes) if ax.name in scan_axes]
    if len(check_list) < len(scan_axes):
        raise ValueError(
            "Axis not found in the list of goniometer axes. Please check your input."
            f"Goniometer axes: {goniometer_axes}. Looking for {scan_axes}."
        )

    # Check that the chip dict has been passed, raise error if not
    if not chip_info:
        logger.error("No chip_dict found.")
        raise ValueError(
            "No information about the FT chip has been passed. \
            Impossible to determine scan parameters. NeXus file won't be written."
        )

    # Define chip
    chip = Chip(
        "fastchip",
        num_steps=[chip_info["X_NUM_STEPS"][1], chip_info["Y_NUM_STEPS"][1]],
        step_size=[chip_info["X_STEP_SIZE"][1], chip_info["Y_STEP_SIZE"][1]],
        num_blocks=[chip_info["X_NUM_BLOCKS"][1], chip_info["Y_NUM_BLOCKS"][1]],
        block_size=[chip_info["X_BLOCK_SIZE"][1], chip_info["Y_BLOCK_SIZE"][1]],
        start_pos=[
            chip_info["X_START"][1],
            chip_info["Y_START"][1],
            chip_info["Z_START"][1],
        ],
    )

    # Read chip map
    blocks = read_chip_map(
        chipmap,
        chip.num_blocks[0],
        chip.num_blocks[1],
    )

    # Workaround for eg. I19 Eiger which saves an increment for phi/omega in meta file.
    for ax in goniometer_axes:
        if ax.transformation_type == "rotation":
            ax.increment = 0.0

    # Calculate scan start positions on chip
    if list(blocks.values())[0] == "fullchip":
        logger.info("Full chip: all the blocks will be scanned.")
        starts = compute_goniometer(chip, full=True, ax1=scan_axes[0], ax2=scan_axes[1])
        starts = fullchip_blocks_conversion(starts, chip)
    else:
        logger.info(f"Scanning blocks: {list(blocks.keys())}.")
        starts = compute_goniometer(chip, blocks, ax1=scan_axes[0], ax2=scan_axes[1])

    # Create two temporary axes to be used for scan calculations
    axis1 = Axis(scan_axes[0], "", TransformationType.TRANSLATION, (0, 0, 0))
    axis2 = Axis(scan_axes[1], "", TransformationType.TRANSLATION, (0, 0, 0))

    # Iterate over blocks to calculate scan points
    SCAN = {axis1.name: np.array([]), axis2.name: np.array([])}
    for k, v in starts.items():
        axis1.start_pos = v[axis1.name]
        axis1.increment = chip.step_size[1] * v["direction"]
        axis1.num_steps = chip.num_steps[1]
        axis2.start_pos = v[axis2.name]
        axis2.increment = chip.step_size[0]
        axis2.num_steps = chip.num_steps[0]

        logger.debug(
            f"Current block: {k}\n"
            f"{axis1.name} start: {v[axis1.name]} \n"
            f"{axis2.name} start: {v[axis2.name]} \n"
            f"Scan direction: {v['direction']} \n"
        )

        _scan = calculate_scan_points(axis1, axis2)
        SCAN[axis1.name] = np.append(SCAN[axis1.name], np.round(_scan[axis1.name], 3))
        SCAN[axis2.name] = np.append(SCAN[axis2.name], np.round(_scan[axis2.name], 3))

    # Check the number of exposures per window
    N = int(chip_info["N_EXPOSURES"][1])
    if N > 1:
        # Repeat each position N times
        SCAN = {k: [val for val in v for _ in range(N)] for k, v in SCAN.items()}
    logger.info(f"Each position has been collected {N} times.")
    logger.info(f"Pump repeat setting: {chip_info['PUMP_REPEAT'][1]}.")
    pump_info = pump_probe.to_dict()
    pump_info["n_exposures"] = N

    return SCAN, pump_info


def run_3D_grid_scan(
    goniometer_axes: List[Axis],
    chip_info: Dict[str, List],
    chipmap: Path | str,
    pump_probe: PumpProbe,
    osc_axis: str = "omega",
) -> Tuple[Dict]:
    """_summary_

    Args:
        goniometer_axes (List[Axis]): _description_
        chip_info (Dict[str, List]): _description_
        chipmap (Path | str): _description_
        pump_probe (PumpProbe): _description_
        osc_axis (str, optional): _description_

    Returns:
        Tuple[Dict]:
            OSC: dictionary with oscillation scan axis values
            TRANSL: dictionary with grid scan values
            pump_info: updated pump probe information
    """
    logger.info("Running a 3D grid scan experiment.")

    N = int(chip_info["N_EXPOSURES"][1])

    pump_info = pump_probe.to_dict()
    pump_info["repeat"] = int(chip_info["PUMP_REPEAT"][1])
    pump_info["n_exposures"] = N
    return None, None, pump_info
