import os
import cv2 as cv
import numpy as np
import pandas as pd
import tqdm
from typing import Tuple, List, Union
from vame.logging.logger import VameLogger, TqdmToLogger
from pathlib import Path
from vame.util.auxiliary import read_config
from vame.schemas.states import EgocentricAlignmentFunctionSchema, save_state
from vame.schemas.project import PoseEstimationFiletype
from vame.util.data_manipulation import (
    interpol_first_rows_nans,
    crop_and_flip,
    background,
    read_pose_estimation_file,
)


logger_config = VameLogger(__name__)
logger = logger_config.logger


def align_mouse(
    project_path: str,
    session: str,
    video_format: str,
    crop_size: Tuple[int, int],
    pose_list: List[np.ndarray],
    pose_ref_index: Tuple[int, int],
    confidence: float,
    pose_flip_ref: Tuple[int, int],
    bg: np.ndarray,
    frame_count: int,
    use_video: bool = True,
    tqdm_stream: Union[TqdmToLogger, None] = None,
) -> Tuple[List[np.ndarray], List[List[np.ndarray]], np.ndarray]:
    """
    Align the mouse in the video frames.

    Parameters:
    -----------
    project_path : str
        Path to the project directory.
    session : str
        Name of the session.
    video_format : str
        Format of the video file.
    crop_size : Tuple[int, int]
        Size to crop the video frames.
    pose_list : List[np.ndarray]
        List of pose coordinates.
    pose_ref_index : Tuple[int, int]
        Pose reference indices.
    confidence : float
        Pose confidence threshold.
    pose_flip_ref : Tuple[int, int]
        Reference indices for flipping.
    bg : np.ndarray
        Background image.
    frame_count : int
        Number of frames to align.
    use_video : bool, optional
        bool if video should be cropped or DLC points only. Defaults to True.
    tqdm_stream : Union[TqdmToLogger, None], optional
        Tqdm stream to log the progress. Defaults to None.

    Returns
    -------
    Tuple[List[np.ndarray], List[List[np.ndarray]], np.ndarray]
        List of aligned images, list of aligned DLC points, and aligned time series data.
    """
    images = []
    points = []
    for i in pose_list:
        for j in i:
            if j[2] <= confidence:
                j[0], j[1] = np.nan, np.nan

    for i in pose_list:
        i = interpol_first_rows_nans(i)

    if use_video:
        video_path = str(
            os.path.join(
                project_path,
                "data",
                "raw",
                session + video_format,
            )
        )
        capture = cv.VideoCapture(video_path)
        if not capture.isOpened():
            raise Exception(f"Unable to open video file: {video_path}")

    for idx in tqdm.tqdm(
        range(frame_count),
        disable=not True,
        file=tqdm_stream,
        desc="Align frames",
    ):
        if use_video:
            # Read frame
            try:
                ret, frame = capture.read()
                frame = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
                frame = frame - bg
                frame[frame <= 0] = 0
            except Exception:
                logger.info("Couldn't find a frame in capture.read(). #Frame: %d" % idx)
                continue
        else:
            frame = np.zeros((1, 1))

        # Read coordinates and add border
        pose_list_bordered = []

        for i in pose_list:
            pose_list_bordered.append(
                (int(i[idx][0] + crop_size[0]), int(i[idx][1] + crop_size[1]))
            )

        img = cv.copyMakeBorder(
            frame,
            crop_size[1],
            crop_size[1],
            crop_size[0],
            crop_size[0],
            cv.BORDER_CONSTANT,
            0,
        )

        coord_center = []
        punkte = []

        for i in pose_ref_index:
            coord = []
            # changed from pose_list_bordered[i][0] 2/28/2024 PN
            coord.append(pose_list_bordered[i][0])
            # changed from pose_list_bordered[i][1] 2/28/2024 PN
            coord.append(pose_list_bordered[i][1])
            punkte.append(coord)

        # coord_center.append(pose_list_bordered[5][0]-5)
        # coord_center.append(pose_list_bordered[5][0]+5)
        # coord_center = [coord_center]
        punkte = [punkte]

        # coord_center = np.asarray(coord_center)
        punkte = np.asarray(punkte)

        # calculate minimal rectangle around snout and tail
        rect = cv.minAreaRect(punkte)
        # rect_belly = cv.minAreaRect(coord_center)
        # center_belly, size_belly, theta_belly = rect_belly
        # change size in rect tuple structure to be equal to crop_size
        lst = list(rect)
        lst[1] = crop_size
        # lst[0] = center_belly
        rect = tuple(lst)
        center, size, theta = rect

        # lst2 = list(rect)
        # lst2[0][0] = center[0] - size[0]//2
        # lst2[0][1] = center[1] - size[1]//2
        # rect = tuple(lst2)
        # center[0] -= size[0]//2
        # center[1] -= size[0]//2 # added this shift to change center to belly 2/28/2024

        # crop image
        out, shifted_points = crop_and_flip(
            rect,
            img,
            pose_list_bordered,
            pose_flip_ref,
        )

        if use_video:  # for memory optimization, just save images when video is used.
            images.append(out)
        points.append(shifted_points)

    if use_video:
        capture.release()

    time_series = np.zeros((len(pose_list) * 2, frame_count))
    for i in range(frame_count):
        idx = 0
        for j in range(len(pose_list)):
            time_series[idx : idx + 2, i] = points[i][j]
            idx += 2

    return images, points, time_series


def play_aligned_video(
    a: List[np.ndarray],
    n: List[List[np.ndarray]],
    frame_count: int,
) -> None:
    """
    Play the aligned video.

    Parameters
    ----------
    a : List[np.ndarray]
        List of aligned images.
    n : List[List[np.ndarray]]
        List of aligned DLC points.
    frame_count : int
        Number of frames in the video.
    """
    colors = [
        (255, 0, 0),
        (0, 255, 0),
        (0, 0, 255),
        (255, 255, 0),
        (255, 0, 255),
        (0, 255, 255),
        (0, 0, 0),
        (255, 255, 255),
    ]
    for i in range(frame_count):
        # Capture frame-by-frame
        ret, frame = True, a[i]
        if ret is True:
            # Display the resulting frame
            frame = cv.cvtColor(frame.astype("uint8") * 255, cv.COLOR_GRAY2BGR)
            im_color = cv.applyColorMap(frame, cv.COLORMAP_JET)
            for c, j in enumerate(n[i]):
                cv.circle(im_color, (j[0], j[1]), 5, colors[c], -1)
            cv.imshow("Frame", im_color)
            # Press Q on keyboard to exit
            # Break the loop
            if cv.waitKey(25) & 0xFF == ord("q"):
                break
        else:
            break
    cv.destroyAllWindows()


def alignment(
    project_path: str,
    session: str,
    pose_ref_index: Tuple[int, int],
    video_format: str,
    crop_size: Tuple[int, int],
    confidence: float,
    pose_estimation_filetype: PoseEstimationFiletype,
    path_to_pose_nwb_series_data: Union[str, None] = None,
    use_video: bool = False,
    check_video: bool = False,
    tqdm_stream: Union[TqdmToLogger, None] = None,
) -> Tuple[np.ndarray, List[np.ndarray]]:
    """
    Perform alignment of egocentric data.

    Parameters:
    -----------
    project_path : str
        Path to the project directory.
    session : str
        Name of the session.
    pose_ref_index : List[int]
        Pose reference indices.
    video_format : str
        Format of the video file.
    crop_size : Tuple[int, int]
        Size to crop the video frames.
    confidence : float
        Pose confidence threshold.
    pose_estimation_filetype : PoseEstimationFiletype
        Pose estimation file type. Can be .csv or .nwb.
    path_to_pose_nwb_series_data : Union[str, None], optional
        Path to the pose series data in nwb files. Defaults to None.
    use_video : bool, optional
        Whether to use video for alignment. Defaults to False.
    check_video : bool, optional
        Whether to check the aligned video. Defaults to False.
    tqdm_stream : Union[TqdmToLogger, None], optional
        Tqdm stream to log the progress. Defaults to None.

    Returns
    -------
    Tuple[np.ndarray, List[np.ndarray]]
        Aligned time series data and list of aligned frames.
    """
    # read out data
    file_path = str(Path(project_path) / "data" / "raw" / f"{session}.nc")
    data, data_mat = read_pose_estimation_file(
        file_path=file_path,
        file_type=pose_estimation_filetype,
        path_to_pose_nwb_series_data=path_to_pose_nwb_series_data,
    )

    # get the coordinates for alignment from data table
    # pose_list dimensions: (num_body_parts, num_frames, 3)
    pose_list = []
    for i in range(int(data_mat.shape[1] / 3)):
        pose_list.append(data_mat[:, i * 3 : (i + 1) * 3])

    # list of reference coordinate indices for alignment
    # 0: snout, 1: forehand_left, 2: forehand_right,
    # 3: hindleft, 4: hindright, 5: tail
    # list of 2 reference coordinate indices for avoiding flipping
    pose_flip_ref = pose_ref_index

    if use_video:
        # compute background
        video_path = str(
            os.path.join(
                project_path,
                "data",
                "raw",
                session + video_format,
            )
        )
        bg = background(
            project_path=project_path,
            session=session,
            video_path=video_path,
            save_background=False,
        )
        capture = cv.VideoCapture(video_path)
        if not capture.isOpened():
            raise Exception(f"Unable to open video file: {video_path}")
        frame_count = int(capture.get(cv.CAP_PROP_FRAME_COUNT))
        capture.release()
    else:
        bg = 0
        # Change this to an abitrary number if you first want to test the code
        frame_count = len(data)

    frames, n, time_series = align_mouse(
        project_path=project_path,
        session=session,
        video_format=video_format,
        crop_size=crop_size,
        pose_list=pose_list,
        pose_ref_index=pose_ref_index,
        confidence=confidence,
        pose_flip_ref=pose_flip_ref,
        bg=bg,
        frame_count=frame_count,
        use_video=use_video,
        tqdm_stream=tqdm_stream,
    )

    if check_video:
        play_aligned_video(frames, n, frame_count)

    return time_series, frames


@save_state(model=EgocentricAlignmentFunctionSchema)
def egocentric_alignment(
    config: str,
    pose_ref_index: Tuple[int, int] = (0, 1),
    crop_size: Tuple[int, int] = (300, 300),
    use_video: bool = False,
    video_format: str = ".mp4",
    check_video: bool = False,
    save_logs: bool = False,
) -> None:
    """
    Egocentric alignment of bevarioral videos.
    Fills in the values in the "egocentric_alignment" key of the states.json file.
    Creates training dataset for VAME at:
    - project_name/
        - data/
            - filename/
                - filename-PE-seq.npy
            - filename/
                - filename-PE-seq.npy
    The produced .npy files contain the aligned time series data in the
    shape of (num_dlc_features, num_video_frames).

    Parameters
    ----------
    config : str
        Path for the project config file.
    pose_ref_index : list, optional
        Pose reference index to be used to align. Defaults to [0, 1].
    crop_size : tuple, optional
        Size to crop the video. Defaults to (300,300).
    use_video : bool, optional
        Weather to use video to do the post alignment. Defaults to False.
    video_format : str, optional
        Video format, can be .mp4 or .avi. Defaults to '.mp4'.
    check_video : bool, optional
        Weather to check the video. Defaults to False.

    Raises:
    ------
    ValueError
        If the config.yaml indicates that the data is not egocentric.
    """
    try:
        config_file = Path(config).resolve()
        cfg = read_config(str(config_file))
        if cfg["egocentric_data"]:
            raise ValueError(
                "The config.yaml indicates that the data is egocentric. Please check the parameter 'egocentric_data'."
            )
        tqdm_stream = None

        if save_logs:
            log_path = Path(cfg["project_path"]) / "logs" / "egocentric_alignment.log"
            logger_config.add_file_handler(str(log_path))
            tqdm_stream = TqdmToLogger(logger=logger)

        logger.info("Starting egocentric alignment")
        project_path = cfg["project_path"]
        sessions = cfg["session_names"]
        confidence = cfg["pose_confidence"]
        num_features = cfg["num_features"]
        video_format = video_format
        crop_size = crop_size

        y_shifted_indices = np.arange(0, num_features, 2)
        x_shifted_indices = np.arange(1, num_features, 2)
        belly_Y_ind = pose_ref_index[0] * 2
        belly_X_ind = (pose_ref_index[0] * 2) + 1

        # call function and save into your VAME data folder
        paths_to_pose_nwb_series_data = cfg["paths_to_pose_nwb_series_data"]
        for i, session in enumerate(sessions):
            logger.info(
                "Aligning session %s, Pose confidence value: %.2f"
                % (session, confidence)
            )
            egocentric_time_series, frames = alignment(
                project_path=project_path,
                session=session,
                pose_ref_index=pose_ref_index,
                video_format=video_format,
                crop_size=crop_size,
                confidence=confidence,
                pose_estimation_filetype=cfg["pose_estimation_filetype"],
                path_to_pose_nwb_series_data=(
                    paths_to_pose_nwb_series_data
                    if not paths_to_pose_nwb_series_data
                    else paths_to_pose_nwb_series_data[i]
                ),
                use_video=use_video,
                check_video=check_video,
                tqdm_stream=tqdm_stream,
            )

            # Shifiting section added 2/29/2024 PN
            # TODO - should this be hardcoded like that?
            egocentric_time_series_shifted = egocentric_time_series
            belly_Y_shift = egocentric_time_series[belly_Y_ind, :]
            belly_X_shift = egocentric_time_series[belly_X_ind, :]

            egocentric_time_series_shifted[y_shifted_indices, :] -= belly_Y_shift
            egocentric_time_series_shifted[x_shifted_indices, :] -= belly_X_shift

            # Save new shifted file
            np.save(
                os.path.join(
                    project_path, "data", "processed", session, session + "-PE-seq.npy"
                ),
                egocentric_time_series_shifted,
            )

        logger.info(
            "Your data is now in the right format and you can call vame.create_trainset()"
        )
    except Exception as e:
        logger.exception(f"{e}")
        raise e
    finally:
        logger_config.remove_file_handler()
