"""
This code is meant to 1) compress the latent representation of the data contained in latent_location.npy 
with shape (n_samples,n_features,n_perturbations) to 3 UMAP dimensions.
Then, the location of the samples in 3D-UMAP is plotted, color coding by a feature of interest. The movement
of the samples when perturbing said figure is shown using the same UMAP projection as the baseline.

Args:
    -lp or --latent_path: path to latent numpy array of shape (n_samples,n_features,n_perturbations)
    -dp or --data_path: path to original datasets, in the example is interim_path
    -ds or --dataset: name of the perturbed dataset
    -foi or --feature_of_interest: feature that we want to perturb and visualize

Returns:
    figure folder inside latent_path/ with figures and gifs depicting the latent space distribution of
    the feature of interest (perturbed_feature.gif) and the movement that samples undergo when perturbing 
    said feature (arrows.gif)

Example: 
    args.latent_path = Path("/Users/_____/Desktop/MOVE/tutorial/results/identify_associations")
    args.data_path = Path("/Users/_____/Desktop/MOVE/tutorial/interim_data")
    args.dataset = "ibd.mbx"
    args.feature_of_interest = "C20 carnitine"

How to run:
    Example:
    1) go to the folder where this file is located:
    cd /Users/____/Desktop/MOVE/supplementary_files
    2) type the following substituting the fields for your files
    python 3D_latent_visualization.py -lp /Users/____/Desktop/MOVE/tutorial/results/identify_associations \\
                                      -dp /Users/____/Desktop/MOVE/tutorial/interim_data \\
                                      -ds ibd.mbx \\
                                      -foi="C20 carnitine"  
                                      
Note: UMAP must be installed, which can be done by running:
    pip install umap-learn
"""

import argparse
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import umap
from PIL import Image

from move.visualization.latent_space import plot_3D_latent_and_displacement

parser = argparse.ArgumentParser(
    description="Read latent space matrix file to plot it in 3D"
)
parser.add_argument(
    "-lp",
    "--latent_path",
    metavar="lp",
    type=Path,
    required=True,
    help="path to latent numpy array (n_samples,n_features,n_perturbations)",
)
parser.add_argument(
    "-dp",
    "--data_path",
    metavar="dp",
    type=Path,
    required=True,
    help="path to original datasets, interim_path",
)
parser.add_argument(
    "-ds",
    "--dataset",
    metavar="ds",
    type=str,
    required=True,
    help="name of the perturbed dataset",
)
parser.add_argument(
    "-foi",
    "--feature_of_interest",
    metavar="foi",
    type=str,
    required=True,
    help="feature that we want to perturb",
)
args = parser.parse_args()

figure_path = Path(args.latent_path / "figures")
figure_path.mkdir(exist_ok=True, parents=True)

perturbed_dataset = np.load(args.data_path / f"{args.dataset}.npy")
perturbed_features = list(np.load(args.latent_path / "perturbed_features_list.npy"))

latent_matrix = np.load(args.latent_path / "latent_location.npy")
trans = umap.UMAP(random_state=42, n_components=3).fit(latent_matrix[:, :, -1])
embedding = trans.embedding_

if args.feature_of_interest not in perturbed_features:
    raise ValueError(" Feature of interest not in perturbed dataset")

i = perturbed_features.index(args.feature_of_interest)

new_embedding = trans.transform(latent_matrix[:, :, i])

# # Plot latent space:
pic_num = 0
n_pictures = 100
for azimuth, altitude in zip(
    np.linspace(-45, 45, n_pictures), np.linspace(15, 45, n_pictures)
):
    fig = plot_3D_latent_and_displacement(
        embedding,
        new_embedding,
        feature_values=perturbed_dataset[:, i],
        feature_name=f"Sample movement",
        show_baseline=False,
        show_perturbed=False,
        show_arrows=True,
        step=1,
        altitude=altitude,
        azimuth=azimuth,
    )

    fig.savefig(figure_path / f"3D_latent_movement_{pic_num}_arrows.png", dpi=100)
    plt.close(fig)

    fig = plot_3D_latent_and_displacement(
        embedding,
        new_embedding,
        feature_values=perturbed_dataset[:, i],
        feature_name=f"Feature {args.feature_of_interest}",
        show_baseline=True,
        show_perturbed=False,
        show_arrows=False,
        altitude=altitude,
        azimuth=azimuth,
    )
    fig.savefig(
        figure_path / f"3D_latent_movement_{pic_num}_perturbed_feature.png", dpi=100
    )
    plt.close(fig)
    pic_num += 1

for plot_type in ["arrows", "perturbed_feature"]:
    frames = [
        Image.open(figure_path / f"3D_latent_movement_{pic_num}_{plot_type}.png")
        for pic_num in range(n_pictures)
    ]  # sorted(glob.glob("*3D_latent*"))]
    frames[0].save(
        figure_path / f"{plot_type}.gif",
        format="GIF",
        append_images=frames[1:],
        save_all=True,
        duration=75,
        loop=0,
    )