#!python3
""" CLI to simplify analysing the output of match_template.py.

    Copyright (c) 2023 European Molecular Biology Laboratory

    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
"""
import argparse
from os.path import splitext

import numpy as np

from tme import Density, Orientations
from tme.matching_utils import (
    create_mask,
    generate_tempfile_name,
    rotation_aligning_vectors,
    euler_from_rotationmatrix,
)


def parse_args():
    parser = argparse.ArgumentParser(
        description="Extract matching candidates for further refinement."
    )

    io_group = parser.add_argument_group("Input / Output")
    io_group.add_argument(
        "--input_file",
        required=False,
        type=str,
        help="Path to the output of match_template.py.",
    )
    io_group.add_argument(
        "--orientations",
        required=True,
        type=str,
        help="Path to file generated by postprocess.py using output_format orientations.",
    )
    io_group.add_argument(
        "--target",
        required=False,
        type=str,
        help="Extract candidates from this target, can be at different sampling rate.",
    )
    io_group.add_argument(
        "--template",
        required=False,
        type=str,
        help="Extract candidates from this target, can be at different sampling rate.",
    )
    io_group.add_argument(
        "-o",
        "--output_file",
        required=False,
        type=str,
        help="Path to output HDF5 file.",
    )

    extraction_group = parser.add_argument_group("Extraction")
    extraction_group.add_argument(
        "--box_size",
        required=False,
        type=int,
        help="Box size for extraction, defaults to two times the template box.",
    )
    extraction_group.add_argument(
        "--alignment_vector",
        required=False,
        type=str,
        help="Vector to align extracted coordinate to e.g. '1,0,0'. If provided, euler_z"
        "euler_y and euler_x are assumed to be a vector that will be rotated to match "
        "alignment_vector.",
    )
    extraction_group.add_argument(
        "--interpolation_order",
        dest="interpolation_order",
        required=False,
        type=int,
        default=1,
        help="Spline interpolation used for template rotations. If less than zero "
        "no interpolation is performed.",
    )
    extraction_group.add_argument(
        "--translation_uncertainty",
        dest="translation_uncertainty",
        required=False,
        type=int,
        default=None,
        help="Creates a cented spherical target mask with given radius in voxel.",
    )
    extraction_group.add_argument(
        "--mask_path",
        dest="mask_path",
        required=False,
        type=str,
        default="target_mask.h5",
        help="Path to write spherical mask to, defaults to target_mask.h5.",
    )
    args = parser.parse_args()

    if args.alignment_vector is not None:
        args.alignment_vector = [float(x) for x in args.alignment_vector.split(",")]
        if len(args.alignment_vector) != 3:
            raise ValueError("Only three-dimensional vectors are supported")

    return args


def main():
    args = parse_args()
    orientations = Orientations.from_file(args.orientations)

    target = Density.from_file(args.target, use_memmap=True)
    template = Density.from_file(args.template)

    box_size = template.shape
    if args.box_size is not None:
        args.box_size = args.box_size // 2
        box_size = np.full(template.data.ndim, fill_value=args.box_size, dtype=int)

    subtomo_shape = np.multiply(box_size, 2).astype(int)
    extraction_scaling = 3 if args.alignment_vector is not None else 2
    extraction_shape = np.multiply(box_size, extraction_scaling).astype(int)
    orientations, cand_slices, obs_slices = orientations.get_extraction_slices(
        target_shape=target.shape,
        extraction_shape=extraction_shape,
        drop_out_of_box=True,
        return_orientations=True,
    )

    filename = generate_tempfile_name()
    output_dtype = target.data.dtype
    if args.alignment_vector is not None:
        output_dtype = np.float32

    if args.translation_uncertainty is not None:
        dens = Density(
            np.memmap(
                args.mask_path,
                mode="w+",
                shape=(len(obs_slices), *subtomo_shape),
                dtype=output_dtype,
            ),
            sampling_rate=(1, *target.sampling_rate),
            origin=(0, *target.origin),
        )
        dens.data[:] = 0
        mask = create_mask(
            mask_type="ellipse",
            center=template.shape,
            radius=args.translation_uncertainty,
            shape=subtomo_shape,
        )
        dens.data[:] = mask
        dens.to_file(args.mask_path)

    target.data = target.data.astype(output_dtype)

    dens = Density(
        np.memmap(
            filename,
            mode="w+",
            shape=(len(obs_slices), *subtomo_shape),
            dtype=output_dtype,
        ),
        sampling_rate=(1, *target.sampling_rate),
        origin=(0, *target.origin),
    )
    dens.data[:] = target.metadata["mean"]

    target_meta = {
        k: v for k, v in target.metadata.items() if k in ("mean", "max", "min", "std")
    }

    total_slices = len(obs_slices)
    magnitude = len(str(total_slices))
    for index, obs_slice in enumerate(obs_slices):
        print(f"Processing {index:{magnitude}} / {total_slices:{magnitude}}", end="\r")

        target_subset = Density(
            target.data[obs_slice],
            sampling_rate=target.sampling_rate,
            origin=target.origin,
        )

        rotation_matrix = np.eye(target_subset.data.ndim)
        if args.alignment_vector is not None:
            normal = orientations.rotations[index]
            # To visualize only normals were picked we would need to do the opposite
            # However, this allows us to align everything to the template axis
            rotation_matrix = rotation_aligning_vectors(
                normal, target_vector=args.alignment_vector
            )

            target_subset = target_subset.rigid_transform(
                rotation_matrix=rotation_matrix,
                use_geometric_center=True,
                order=args.interpolation_order,
            )
        target_subset.pad(subtomo_shape, center=True)
        dens.data[index] = target_subset.data
        orientations.rotations[index] = euler_from_rotationmatrix(rotation_matrix)

    dens.metadata.update(target_meta)
    dens.to_file(args.output_file)
    orientations.to_file(
        f"{splitext(args.output_file)}_aligned.tsv",
        file_format="text"
    )
    orientations.to_file(
        f"{splitext(args.output_file)}_aligned.star",
        file_format="relion"
    )


if __name__ == "__main__":
    main()
