from __future__ import annotations

import os
import re
import string
import subprocess
from collections import ChainMap
from math import hypot
from pathlib import Path
from typing import Optional

import plotly.express as px
import workflows.recipe
from gemmi import cif
from pydantic import BaseModel, Field, ValidationError, validator
from workflows.services.common_service import CommonService

from cryoemservices.util.spa_relion_service_options import (
    RelionServiceOptions,
    update_relion_options,
)


class MotionCorrParameters(BaseModel):
    movie: str = Field(..., min_length=1)
    mrc_out: str = Field(..., min_length=1)
    experiment_type: str
    pixel_size: float
    dose_per_frame: float
    use_motioncor2: bool = True
    patch_sizes: dict = {"x": 5, "y": 5}
    gpu: int = 0
    threads: int = 1
    gain_ref: str = None
    rot_gain: int = None
    flip_gain: int = None
    dark: str = None
    use_gpus: int = None
    sum_range: Optional[dict] = None
    iter: int = None
    tol: float = None
    throw: int = None
    trunc: int = None
    fm_ref: int = 0
    voltage: int = None
    fm_int_file: str = None
    init_dose: float = None
    mag: Optional[dict] = None
    motion_corr_binning: float = None
    serial: int = None
    in_suffix: str = None
    eer_sampling: int = None
    out_stack: int = None
    bft: Optional[dict] = None
    group: int = None
    defect_file: str = None
    arc_dir: str = None
    in_fm_motion: int = None
    split_sum: int = None
    dose_motionstats_cutoff: float = 4.0
    do_icebreaker_jobs: bool = True
    movie_id: int
    mc_uuid: int
    picker_uuid: int
    relion_options: Optional[RelionServiceOptions] = None
    ctf: dict = {}

    @validator("experiment_type")
    def is_spa_or_tomo(cls, experiment):
        if experiment not in ["spa", "tomography"]:
            raise ValueError("Specify an experiment type of spa or tomography.")
        return experiment

    class Config:
        ignore_extra = True


class ChainMapWithReplacement(ChainMap):
    def __init__(self, *maps, substitutions=None) -> None:
        super().__init__(*maps)
        self._substitutions = substitutions

    def __getitem__(self, k):
        v = super().__getitem__(k)
        if self._substitutions and isinstance(v, str) and "$" in v:
            template = string.Template(v)
            return template.substitute(**self._substitutions)
        return v


class MotionCorr(CommonService):
    """
    A service for motion correcting cryoEM movies using MotionCor2
    """

    # Human readable service name
    _service_name = "MotionCorr"

    # Logger name
    _logger_name = "cryoemservices.services.motioncorr"

    # Job name
    job_type = "relion.motioncorr"

    # Values to extract for ISPyB
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.x_shift_list = []
        self.y_shift_list = []

    def initializing(self):
        """Subscribe to a queue. Received messages must be acknowledged."""
        self.log.info("Motion correction service starting")
        workflows.recipe.wrap_subscribe(
            self._transport,
            "motioncorr",
            self.motion_correction,
            acknowledgement=True,
            log_extender=self.extend_log,
            allow_non_recipe_messages=True,
        )

    def parse_mc2_stdout(self, mc_stdout: str):
        """
        Read the output logs of MotionCor2 to determine
        the movement of each frame
        """
        frames_line = False
        for line in mc_stdout.split("\n"):
            # Frame reading in MotionCorr 1.4.0
            if line.startswith("...... Frame"):
                line_split = line.split()
                self.x_shift_list.append(float(line_split[-2]))
                self.y_shift_list.append(float(line_split[-1]))

            # Alternative frame reading for MotionCorr 1.6.3
            if not line:
                frames_line = False
            if frames_line:
                line_split = line.split()
                self.x_shift_list.append(float(line_split[1]))
                self.y_shift_list.append(float(line_split[2]))
            if "x Shift" in line:
                frames_line = True

    def parse_relion_mc_output(self, stdout_file):
        """
        Read the output star file made by Relion to determine
        the movement of each frame
        """
        shift_cif = cif.read_file(str(stdout_file))
        shift_block = shift_cif.find_block("global_shift")
        x_shifts_str = list(shift_block.find_loop("_rlnMicrographShiftX"))
        y_shifts_str = list(shift_block.find_loop("_rlnMicrographShiftY"))
        for frame in range(len(x_shifts_str)):
            self.x_shift_list.append(float(x_shifts_str[frame]))
            self.y_shift_list.append(float(y_shifts_str[frame]))

    def motioncor2(self, command, mrc_out):
        """Run the MotionCor2 command"""
        result = subprocess.run(command, capture_output=True)
        self.parse_mc2_stdout(result.stdout.decode("utf8", "replace"))
        return result

    def relion_motioncorr(self, command, mrc_out):
        """Run Relion's owm motion correction"""
        result = subprocess.run(command, capture_output=True)
        if Path(mrc_out).with_suffix(".star").exists():
            self.parse_relion_mc_output(Path(mrc_out).with_suffix(".star"))
        else:
            self.log.error(
                f"Relion output log {Path(mrc_out).with_suffix('.star')} not found"
            )
        return result

    def motion_correction(self, rw, header: dict, message: dict):
        class MockRW:
            def dummy(self, *args, **kwargs):
                pass

        if not rw:
            if (
                not isinstance(message, dict)
                or not message.get("parameters")
                or not message.get("content")
            ):
                self.log.error("Rejected invalid simple message")
                self._transport.nack(header)
                return

            # Create a wrapper-like object that can be passed to functions
            # as if a recipe wrapper was present.
            rw = MockRW()
            rw.transport = self._transport
            rw.recipe_step = {"parameters": message["parameters"], "output": None}
            rw.environment = {"has_recipe_wrapper": False}
            rw.set_default_channel = rw.dummy
            rw.send = rw.dummy
            message = message["content"]

        parameter_map = ChainMapWithReplacement(
            message if isinstance(message, dict) else {},
            rw.recipe_step["parameters"],
            substitutions=rw.environment,
        )

        try:
            if isinstance(message, dict):
                mc_params = MotionCorrParameters(**{**dict(parameter_map), **message})
            else:
                mc_params = MotionCorrParameters(**{**dict(parameter_map)})
        except (ValidationError, TypeError) as e:
            self.log.warning(
                f"Motion correction parameter validation failed for message: {message} "
                f"and recipe parameters: {rw.recipe_step.get('parameters', {})} "
                f"with exception: {e}"
            )
            rw.transport.nack(header)
            return

        # Catch any cases where the movie does not exist
        if not Path(mc_params.movie).is_file():
            self.log.warning(f"Movie {mc_params.movie} does not exist")
            rw.transport.nack(header)
            return

        # Generate the plotting paths
        drift_plot_name = str(Path(mc_params.movie).stem) + "_drift_plot.json"
        plot_path = Path(mc_params.mrc_out).parent / drift_plot_name

        # Check if this file has been run before
        if Path(mc_params.mrc_out).is_file():
            job_is_rerun = True
        else:
            job_is_rerun = False

        # Get the eer grouping out of the fractionation file
        eer_grouping = 0
        if (
            mc_params.movie.endswith(".eer")
            and mc_params.fm_int_file
            and Path(mc_params.fm_int_file).is_file()
        ):
            with open(mc_params.fm_int_file, "r") as eer_file:
                eer_values = eer_file.readline()
                try:
                    eer_grouping = eer_values.split()[1]
                except ValueError:
                    self.log.warning("Cannot read eer grouping")

        if mc_params.experiment_type == "spa":
            # Update the relion options
            mc_params.relion_options = update_relion_options(
                mc_params.relion_options, dict(mc_params)
            )
            mc_params.relion_options.eer_grouping = eer_grouping

        # Determine the input and output files
        self.log.info(f"Input: {mc_params.movie} Output: {mc_params.mrc_out}")
        if not Path(mc_params.mrc_out).parent.exists():
            Path(mc_params.mrc_out).parent.mkdir(parents=True)
        if mc_params.movie.endswith(".mrc"):
            input_flag = "-InMrc"
        elif mc_params.movie.endswith((".tif", ".tiff")):
            input_flag = "-InTiff"
        elif mc_params.movie.endswith(".eer"):
            input_flag = "-InEer"
        else:
            self.log.error(f"No input flag found for movie {mc_params.movie}")
            rw.transport.nack(header)
            return

        # Run motion correction
        if mc_params.use_motioncor2:
            # Construct the command for MotionCor2
            self.job_type = "relion.motioncorr.motioncor2"
            self.log.info("Using MotionCor2")
            command = ["MotionCor2", input_flag, mc_params.movie]
            mc2_flags = {
                "mrc_out": "-OutMrc",
                "patch_sizes": "-Patch",
                "pixel_size": "-PixSize",
                "gain_ref": "-Gain",
                "rot_gain": "-RotGain",
                "flip_gain": "-FlipGain",
                "dark": "-Dark",
                "gpu": "-Gpu",
                "use_gpus": "-UseGpus",
                "sum_range": "-SumRange",
                "iter": "-Iter",
                "tol": "-Tol",
                "throw": "-Throw",
                "trunc": "-Trunc",
                "fm_ref": "-FmRef",
                "voltage": "-Kv",
                "dose_per_frame": "-FmDose",
                "init_dose": "-InitDose",
                "mag": "-Mag",
                "motion_corr_binning": "-FtBin",
                "serial": "-Serial",
                "in_suffix": "-InSuffix",
                "eer_sampling": "-EerSampling",
                "out_stack": "-OutStack",
                "bft": "-Bft",
                "group": "-Group",
                "defect_file": "-DefectFile",
                "arc_dir": "-ArcDir",
                "in_fm_motion": "-InFmMotion",
                "split_sum": "-SplitSum",
            }
            if mc_params.movie.endswith(".eer"):
                mc2_flags["fm_int_file"] = "-FmIntFile"

            # Add values from input parameters with flags
            for k, v in mc_params.dict().items():
                if (v is not None) and (k in mc2_flags):
                    if type(v) is dict:
                        command.extend(
                            (mc2_flags[k], " ".join(str(_) for _ in v.values()))
                        )
                    else:
                        command.extend((mc2_flags[k], str(v)))
            # Run MotionCor2
            result = self.motioncor2(command, mc_params.mrc_out)

            dose_weighted = Path(mc_params.mrc_out).parent / (
                Path(mc_params.mrc_out).stem + "_DW.mrc"
            )
            if dose_weighted.is_file():
                mc_params.mrc_out = str(dose_weighted)
        else:
            # Construct the command for Relion motion correction
            self.job_type = "relion.motioncorr.own"
            self.log.info("Using Relion's own motion correction")
            os.environ["FI_PROVIDER"] = "tcp"
            command = ["relion_motion_correction", "--use_own"]
            relion_mc_flags = {
                "threads": "--j",
                "movie": "--in_movie",
                "mrc_out": "--out_mic",
                "pixel_size": "--angpix",
                "voltage": "--voltage",
                "dose_per_frame": "--dose_per_frame",
                "gain_ref": "--gainref",
                "defect_file": "--defect_file",
                "rot_gain": "--gain_rot",
                "flip_gain": "--gain_flip",
                "motion_corr_binning": "--bin_factor",
                "eer_sampling": "--eer_upsampling",
                "init_dose": "--preexposure",
                "patch_sizes": {"--patch_x": "x", "--patch_y": "y"},
                "bft": {"--bfactor": "local_motion"},
            }
            # Add values from input parameters with flags
            for param_k, param_v in mc_params.dict().items():
                if (param_v is not None) and (param_k in relion_mc_flags):
                    if type(param_v) is dict:
                        for flag_k, flag_v in relion_mc_flags[param_k].items():
                            command.extend(
                                (flag_k, str(mc_params.dict()[param_k][flag_v]))
                            )
                    else:
                        command.extend((relion_mc_flags[param_k], str(param_v)))

            # Add eer grouping if file is eer
            if eer_grouping:
                command.extend(("--eer_grouping", eer_grouping))

            # Add some standard flags
            command.extend(("--dose_weighting", "--i", "dummy"))
            # Run Relion motion correction
            result = self.relion_motioncorr(command, mc_params.mrc_out)

        # Adjust the pixel size based on the binning
        if mc_params.motion_corr_binning:
            mc_params.pixel_size *= mc_params.motion_corr_binning
            if mc_params.relion_options:
                mc_params.relion_options.pixel_size *= mc_params.motion_corr_binning

        # Confirm the command ran successfully
        if result.returncode:
            self.log.error(
                f"Motion correction of {mc_params.movie} "
                f"failed with exitcode {result.returncode}:\n"
                + result.stderr.decode("utf8", "replace")
            )
            if mc_params.experiment_type == "spa" and not job_is_rerun:
                # On spa failure send the outputs to the node creator
                node_creator_parameters = {
                    "job_type": self.job_type,
                    "input_file": mc_params.mrc_out,
                    "output_file": mc_params.mrc_out,
                    "relion_options": dict(mc_params.relion_options),
                    "command": " ".join(command),
                    "stdout": result.stdout.decode("utf8", "replace"),
                    "stderr": result.stderr.decode("utf8", "replace"),
                    "success": False,
                }
                if isinstance(rw, MockRW):
                    rw.transport.send(
                        destination="node_creator",
                        message={
                            "parameters": node_creator_parameters,
                            "content": "dummy",
                        },
                    )
                else:
                    rw.send_to("node_creator", node_creator_parameters)
            rw.transport.nack(header)
            return

        # Extract results for ispyb
        total_motion = 0
        early_motion = 0
        late_motion = 0
        cutoff_frame = round(
            mc_params.dose_motionstats_cutoff / mc_params.dose_per_frame
        )
        for i in range(1, len(self.x_shift_list)):
            total_motion += hypot(
                self.x_shift_list[i] - self.x_shift_list[i - 1],
                self.y_shift_list[i] - self.y_shift_list[i - 1],
            )
            if i < cutoff_frame:
                early_motion += hypot(
                    self.x_shift_list[i] - self.x_shift_list[i - 1],
                    self.y_shift_list[i] - self.y_shift_list[i - 1],
                )
            else:
                late_motion += hypot(
                    self.x_shift_list[i] - self.x_shift_list[i - 1],
                    self.y_shift_list[i] - self.y_shift_list[i - 1],
                )
        average_motion_per_frame = total_motion / len(self.x_shift_list)

        # Extract results for ispyb
        fig = px.scatter(x=self.x_shift_list, y=self.y_shift_list)
        fig.write_json(plot_path)
        snapshot_path = Path(mc_params.mrc_out).with_suffix(".jpeg")

        # Forward results to ISPyB
        ispyb_parameters = {
            "ispyb_command": "buffer",
            "buffer_command": {"ispyb_command": "insert_motion_correction"},
            "buffer_store": mc_params.mc_uuid,
            "first_frame": 1,
            "last_frame": len(self.x_shift_list),
            "total_motion": total_motion,
            "average_motion_per_frame": average_motion_per_frame,
            "drift_plot_full_path": str(plot_path),
            "micrograph_snapshot_full_path": str(snapshot_path),
            "micrograph_full_path": str(mc_params.mrc_out),
            "patches_used_x": mc_params.patch_sizes["x"],
            "patches_used_y": mc_params.patch_sizes["y"],
            "dose_per_frame": mc_params.dose_per_frame,
        }
        self.log.info(f"Sending to ispyb {ispyb_parameters}")
        if isinstance(rw, MockRW):
            rw.transport.send(
                destination="ispyb_connector",
                message={
                    "parameters": ispyb_parameters,
                    "content": {"dummy": "dummy"},
                },
            )
        else:
            rw.send_to("ispyb_connector", ispyb_parameters)

        # If this is SPA, determine and set up the next jobs
        if mc_params.experiment_type == "spa":
            # Set up icebreaker if requested, then ctffind
            icebreaker_output = Path(
                re.sub(
                    "MotionCorr/job002/",
                    "IceBreaker/job003/",
                    mc_params.mrc_out,
                )
            )
            if mc_params.do_icebreaker_jobs and not icebreaker_output.is_file():
                # Three IceBreaker jobs: CtfFind job is MC+4
                ctf_job_number = 6

                # Both IceBreaker micrographs and flattening inherit from motioncorr
                self.log.info(
                    f"Sending to IceBreaker micrograph analysis: {mc_params.mrc_out}"
                )
                icebreaker_job003_params = {
                    "icebreaker_type": "micrographs",
                    "input_micrographs": mc_params.mrc_out,
                    "output_path": re.sub(
                        "MotionCorr/job002/.+",
                        "IceBreaker/job003/",
                        mc_params.mrc_out,
                    ),
                    "mc_uuid": mc_params.mc_uuid,
                    "relion_options": dict(mc_params.relion_options),
                    "total_motion": total_motion,
                    "early_motion": early_motion,
                    "late_motion": late_motion,
                }
                if isinstance(rw, MockRW):
                    rw.transport.send(
                        destination="icebreaker",
                        message={
                            "parameters": icebreaker_job003_params,
                            "content": "dummy",
                        },
                    )
                else:
                    rw.send_to("icebreaker", icebreaker_job003_params)

                self.log.info(
                    f"Sending to IceBreaker contrast enhancement: {mc_params.mrc_out}"
                )
                icebreaker_job004_params = {
                    "icebreaker_type": "enhancecontrast",
                    "input_micrographs": mc_params.mrc_out,
                    "output_path": re.sub(
                        "MotionCorr/job002/.+",
                        "IceBreaker/job004/",
                        mc_params.mrc_out,
                    ),
                    "mc_uuid": mc_params.mc_uuid,
                    "relion_options": dict(mc_params.relion_options),
                    "total_motion": total_motion,
                    "early_motion": early_motion,
                    "late_motion": late_motion,
                }
                if isinstance(rw, MockRW):
                    rw.transport.send(
                        destination="icebreaker",
                        message={
                            "parameters": icebreaker_job004_params,
                            "content": "dummy",
                        },
                    )
                else:
                    rw.send_to("icebreaker", icebreaker_job004_params)
            elif mc_params.do_icebreaker_jobs and icebreaker_output.is_file():
                # On a rerun, skip IceBreaker jobs but mark the CtfFind job as MC+4
                ctf_job_number = 6
            else:
                # No IceBreaker jobs: CtfFind job is MC+1
                ctf_job_number = 3
            mc_params.ctf["output_image"] = str(
                Path(
                    mc_params.mrc_out.replace(
                        "MotionCorr/job002", f"CtfFind/job{ctf_job_number:03}"
                    )
                ).with_suffix(".ctf")
            )
            mc_params.ctf["relion_options"] = dict(mc_params.relion_options)
            mc_params.ctf["amplitude_contrast"] = mc_params.relion_options.ampl_contrast
        else:
            mc_params.ctf["output_image"] = mc_params.mrc_out.replace(
                "MotionCorr", "CTF"
            ).replace("_motion_corrected", "_ctf")

        # Forward results to ctffind (in both SPA and tomography)
        self.log.info(f"Sending to ctf: {mc_params.mrc_out}")
        mc_params.ctf["experiment_type"] = mc_params.experiment_type
        mc_params.ctf["input_image"] = mc_params.mrc_out
        mc_params.ctf["mc_uuid"] = mc_params.mc_uuid
        mc_params.ctf["picker_uuid"] = mc_params.picker_uuid
        mc_params.ctf["pixel_size"] = mc_params.pixel_size
        if isinstance(rw, MockRW):
            rw.transport.send(  # type: ignore
                destination="ctffind",
                message={"parameters": mc_params.ctf, "content": "dummy"},
            )
        else:
            rw.send_to("ctffind", mc_params.ctf)

        # Forward results to images service
        self.log.info(f"Sending to images service {mc_params.mrc_out}")
        if isinstance(rw, MockRW):
            rw.transport.send(
                destination="images",
                message={
                    "image_command": "mrc_to_jpeg",
                    "file": mc_params.mrc_out,
                },
            )
        else:
            rw.send_to(
                "images",
                {
                    "image_command": "mrc_to_jpeg",
                    "file": mc_params.mrc_out,
                },
            )

        # If this is a new SPA run, send the results to be processed by the node creator
        if mc_params.experiment_type == "spa" and not job_is_rerun:
            # As this is the entry point we need to import the file to the project
            self.log.info("Sending relion.import.movies to node creator")
            project_dir = Path(
                re.search(".+/job[0-9]{3}/", mc_params.mrc_out)[0]
            ).parent.parent
            import_movie = (
                project_dir
                / "Import/job001"
                / Path(mc_params.mrc_out)
                .relative_to(project_dir / "MotionCorr/job002")
                .parent
                / Path(mc_params.movie).name
            )
            if not import_movie.parent.is_dir():
                import_movie.parent.mkdir(parents=True)
            import_movie.unlink(missing_ok=True)
            import_movie.symlink_to(mc_params.movie)
            import_parameters = {
                "job_type": "relion.import.movies",
                "input_file": str(mc_params.movie),
                "output_file": str(import_movie),
                "relion_options": dict(mc_params.relion_options),
                "command": "",
                "stdout": "",
                "stderr": "",
            }
            if isinstance(rw, MockRW):
                rw.transport.send(
                    destination="node_creator",
                    message={"parameters": import_parameters, "content": "dummy"},
                )
            else:
                rw.send_to("node_creator", import_parameters)

            # Then register the motion correction job with the node creator
            self.log.info(f"Sending {self.job_type} to node creator")
            node_creator_parameters = {
                "job_type": self.job_type,
                "input_file": str(import_movie),
                "output_file": mc_params.mrc_out,
                "relion_options": dict(mc_params.relion_options),
                "command": " ".join(command),
                "stdout": result.stdout.decode("utf8", "replace"),
                "stderr": result.stderr.decode("utf8", "replace"),
                "results": {
                    "total_motion": total_motion,
                    "early_motion": early_motion,
                    "late_motion": late_motion,
                },
            }
            if isinstance(rw, MockRW):
                rw.transport.send(
                    destination="node_creator",
                    message={"parameters": node_creator_parameters, "content": "dummy"},
                )
            else:
                rw.send_to("node_creator", node_creator_parameters)

        # Register completion with Murfey if this is tomography
        if mc_params.experiment_type == "tomography":
            self.log.info("Sending to Murfey")
            if isinstance(rw, MockRW):
                rw.transport.send(
                    "murfey_feedback",
                    {
                        "register": "motion_corrected",
                        "movie": mc_params.movie,
                        "mrc_out": mc_params.mrc_out,
                        "movie_id": mc_params.movie_id,
                    },
                )
            else:
                rw.send_to(
                    "murfey_feedback",
                    {
                        "register": "motion_corrected",
                        "movie": mc_params.movie,
                        "mrc_out": mc_params.mrc_out,
                        "movie_id": mc_params.movie_id,
                    },
                )

        self.log.info(f"Done {self.job_type} for {mc_params.movie}.")
        rw.transport.ack(header)
        self.x_shift_list = []
        self.y_shift_list = []
