#!python3
""" Prepare orientations stack for refinement.

    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 (
    load_pickle,
    generate_tempfile_name,
    rotation_aligning_vectors,
    euler_from_rotationmatrix,
    euler_to_rotationmatrix,
)


class ProgressBar:
    """
    ASCII progress bar.
    """

    def __init__(self, message : str, nchars : int, total : int):
        self._size = nchars - len(message)
        self._message = message
        self._total = total

    def update(self, cur):
        x = int(cur * self._size / self._total)
        print(
            "%s[%s%s] %i/%i\r"
            % (self._message, "#" * x, "." * (self._size - x), cur, self._total),
            end=''
        )


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=True,
        type=str,
        help="Path to output HDF5 file.",
    )

    alignment_group = parser.add_argument_group("Alignment")
    alignment_group.add_argument(
        "--align_orientations",
        action="store_true",
        required=False,
        help="Whether to align extracted orientations based on their angles. Allows "
        "for efficient subsequent sampling of cone angles.",
    )
    alignment_group.add_argument(
        "--angles_are_vector",
        action="store_true",
        required=False,
        help="Considers euler_z euler_y, euler_x as vector that will be rotated to align "
        "with the z-axis (1,0,0). Only considered when --align_orientations is set."
    )
    alignment_group.add_argument(
        "--interpolation_order",
        dest="interpolation_order",
        required=False,
        type=int,
        default=1,
        help="Interpolation order for alignment, less than zero is no interpolation."
    )

    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.",
    )
    extraction_group.add_argument(
        "--keep_out_of_box",
        action="store_true",
        required=False,
        help="Whether to keep orientations that fall outside the box. If the "
        "orientations are sensible, it is safe to pass this flag.",
    )

    args = parser.parse_args()

    data_present = args.target is not None and args.template is not None
    if args.input_file is None and not data_present:
        raise ValueError(
            "Either --input_file or --target and --template need to be specified."
        )
    elif args.input_file is not None and data_present:
        raise ValueError(
            "Please specific either --input_file or --target and --template."
        )

    return args


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

    if args.input_file is not None:
        data = load_pickle(args.input_file)
        target_origin, _, sampling_rate, cli_args = data[-1]
        args.target, args.template  = cli_args.target, cli_args.template

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

    try:
        template = Density.from_file(args.template)
    except Exception:
        template = Density.from_structure(args.template, sampling_rate = target.sampling_rate)

    box_size = args.box_size
    if box_size is None:
        box_size = np.multiply(template.shape, 2)
    box_size = np.array(box_size)
    box_size = np.repeat(box_size, template.data.ndim // box_size.size).astype(int)

    extraction_shape = np.copy(box_size)
    if args.align_orientations:
        extraction_shape[:] = int(np.linalg.norm(box_size) + 1)

    orientations, cand_slices, obs_slices = orientations.get_extraction_slices(
        target_shape=target.shape,
        extraction_shape=extraction_shape,
        drop_out_of_box=not args.keep_out_of_box,
        return_orientations=True,
    )

    if args.align_orientations:
        orientations.rotations = orientations.rotations.astype(np.float32)
        for index in range(orientations.rotations.shape[0]):
            rotation_matrix = euler_to_rotationmatrix(orientations.rotations[index])
            rotation_matrix = np.linalg.inv(rotation_matrix)
            if args.angles_are_vector:
                rotation_matrix = rotation_aligning_vectors(
                    orientations.rotations[index], target_vector=(1,0,0)
                )
            orientations.rotations[index] = euler_from_rotationmatrix(rotation_matrix)


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

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

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

    print(target.data.shape)
    # There appears to be an isseu with the stack creation. Trace this further
    data_subset = np.zeros(extraction_shape, dtype = target.data.dtype)
    pbar = ProgressBar(message = "Orientation ", nchars = 80, total = len(obs_slices))
    for index, (obs_slice, cand_slice) in enumerate(zip(obs_slices, cand_slices)):
        pbar.update(index + 1)

        data_subset.fill(0)
        data_subset[cand_slice] = target.data[obs_slice]
        target_subset = Density(
            data_subset,
            sampling_rate=target.sampling_rate,
            origin=target.origin,
        )

        if args.align_orientations:
            rotation_matrix = euler_to_rotationmatrix(orientations.rotations[index])
            target_subset = target_subset.rigid_transform(
                rotation_matrix=rotation_matrix,
                use_geometric_center=True,
                order=args.interpolation_order,
            )
        target_subset.pad(box_size, center=True)

        target_value = target.data[tuple(orientations.translations[index].astype(int))]
        center = np.divide(target_subset.data.shape, 2).astype(int ) + np.mod(target_subset.shape, 2)
        print(np.where(target_subset.data == target_value), center)
        print(target_subset.data[tuple(center.astype(int))],
            target_value,
            target_subset.data[tuple(center.astype(int))] == target_value
        )

        dens.data[index] = target_subset.data
    print("")

    target_meta = {
        k: v for k, v in target.metadata.items() if k in ("mean", "max", "min", "std")
    }
    dens.metadata.update(target_meta)
    dens.metadata["batch_dimension"] = (0, )

    dens.to_file(args.output_file)
    orientations.to_file(
        f"{splitext(args.output_file)[0]}_aligned.tsv",
        file_format="text"
    )
    orientations.to_file(
        f"{splitext(args.output_file)[0]}_aligned.star",
        file_format="relion"
    )

if __name__ == "__main__":
    main()
