import argparse
import shutil
from datetime import datetime
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from matplotlib import colormaps
from matplotlib.cm import ScalarMappable
from matplotlib.colors import Colormap, ListedColormap
from mpl_toolkits.axes_grid1.axes_divider import make_axes_locatable
from PIL import Image  # new import


def get_colormap_with_special_missing_color(
    colormap_name: str, missing_data_color: str = "black", reverse: bool = False
) -> Colormap:
    colormap = (
        colormaps[colormap_name] if not reverse else colormaps[f"{colormap_name}_r"]
    )
    colormap.set_bad(missing_data_color)
    colormap.set_over("white")
    return colormap


# --- Setup output folders (similar to data_missing_points.py) ---
# Change the output path as needed
output_path = Path("/home/fedex/mt/plots/data_2d_projections")
datetime_folder_name = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
output_datetime_path = output_path / datetime_folder_name
latest_folder_path = output_path / "latest"
archive_folder_path = output_path / "archive"

for folder in (
    output_path,
    output_datetime_path,
    latest_folder_path,
    archive_folder_path,
):
    folder.mkdir(exist_ok=True, parents=True)

# --- Parse command-line arguments ---
parser = argparse.ArgumentParser(
    description="Plot two 2D projections (from .npy files) in vertical subplots"
)
parser.add_argument(
    "--input1",
    type=Path,
    default=Path(
        "/home/fedex/mt/data/subter/1_loop_closure_illuminated_2023-01-23.npy"
    ),
    help="Path to first .npy file containing 2D projection data",
)
parser.add_argument(
    "--input2",
    type=Path,
    default=Path("/home/fedex/mt/data/subter/3_smoke_human_walking_2023-01-23.npy"),
    help="Path to second .npy file containing 2D projection data",
)
parser.add_argument(
    "--frame1",
    type=int,
    default=955,
    help="Frame index to plot from first file (0-indexed)",
)
parser.add_argument(
    "--frame2",
    type=int,
    default=242,
    help="Frame index to plot from second file (0-indexed)",
)
parser.add_argument(
    "--colormap",
    default="viridis",
    type=str,
    help="Name of matplotlib colormap to use",
)
parser.add_argument(
    "--missing-data-color",
    default="black",
    type=str,
    help="Color to use for missing data in projection",
)
parser.add_argument(
    "--reverse-colormap",
    action="store_true",
    help="Reverse the colormap if specified",
)
args = parser.parse_args()

# --- Load the numpy projection data from the provided files ---
# Each file is assumed to be a 3D array: (num_frames, height, width)
proj_data1 = np.load(args.input1)
proj_data2 = np.load(args.input2)

# Choose the desired frames
try:
    frame1 = proj_data1[args.frame1]
except IndexError:
    raise ValueError(f"Frame index {args.frame1} out of range for file: {args.input1}")
try:
    frame2 = proj_data2[args.frame2]
except IndexError:
    raise ValueError(f"Frame index {args.frame2} out of range for file: {args.input2}")

# Determine shared range across both frames (ignoring NaNs)
global_vmin = 0  # min(np.nanmin(frame1), np.nanmin(frame2))
global_vmax = 0.8  # max(np.nanmax(frame1), np.nanmax(frame2))

# Create colormap using the utility (to mimic create_2d_projection)
cmap = get_colormap_with_special_missing_color(
    args.colormap, args.missing_data_color, args.reverse_colormap
)

# --- Create a figure with 2 vertical subplots and move titles to the left ---
fig, (ax1, ax2) = plt.subplots(nrows=2, ncols=1, figsize=(10, 5))
# leave extra left margin for the left-side labels
fig.subplots_adjust(left=0.14, hspace=0.05)

for ax, frame, label in zip(
    (ax1, ax2),
    (frame1, frame2),
    ("(a)", "(b)"),
):
    im = ax.imshow(frame, cmap=cmap, aspect="auto", vmin=global_vmin, vmax=global_vmax)
    # place the "title" to the left, vertically centered relative to the axes
    ax.text(
        -0.02,  # negative x places text left of the axes (in axes coordinates)
        0.5,
        label,
        transform=ax.transAxes,
        va="center",
        ha="right",
        fontsize=12,
    )
    ax.axis("off")

# Adjust layout to fit margins for a paper
plt.tight_layout(rect=[0, 0.05, 1, 1])
# Add a colorbar with the colormap below the subplots
cbar = fig.colorbar(im, ax=[ax1, ax2], orientation="vertical", fraction=0.05)
cbar.set_label("Reciprocal Range")

# Add a separate colorbar for NaN values
sm = ScalarMappable(cmap=ListedColormap([cmap.get_bad(), cmap.get_over()]))
divider = make_axes_locatable(cbar.ax)
nan_ax = divider.append_axes(
    "bottom", size="15%", pad="3%", aspect=3, anchor=cbar.ax.get_anchor()
)
nan_ax.grid(visible=False, which="both", axis="both")
nan_cbar = fig.colorbar(sm, cax=nan_ax, orientation="vertical")
nan_cbar.set_ticks([0.3, 0.7])
nan_cbar.set_ticklabels(["NaN", f"> {global_vmax}"])
nan_cbar.ax.tick_params(length=0)

# Save the combined plot
output_file = output_datetime_path / "data_2d_projections.png"
plt.savefig(output_file, dpi=300, bbox_inches="tight", pad_inches=0.1)
plt.close()

print(f"Plot saved to: {output_file}")

# --- Create grayscale images (high precision) from the numpy frames using Pillow ---
# Convert NaN values to 0 and ensure the array is in float32 for 32-bit precision.
for degradation_status, frame_number, frame in (
    ("normal", args.frame1, frame1),
    ("smoke", args.frame2, frame2),
):
    frame_gray = np.nan_to_num(frame, nan=0).astype(np.float32)
    gray_image = Image.fromarray(frame_gray, mode="F")
    gray_output_file = (
        output_datetime_path
        / f"frame_{frame_number}_grayscale_{degradation_status}.tiff"
    )
    gray_image.save(gray_output_file)
    print(f"Grayscale image saved to: {gray_output_file}")

# --- Handle folder structure: update latest folder and archive the output folder ---

# Delete current latest folder and recreate it
shutil.rmtree(latest_folder_path, ignore_errors=True)
latest_folder_path.mkdir(exist_ok=True, parents=True)

# Copy contents of the current output datetime folder to latest folder
for file in output_datetime_path.iterdir():
    shutil.copy2(file, latest_folder_path)

# Copy this python script to both output datetime folder and latest folder for preservation
script_path = Path(__file__)
shutil.copy2(script_path, output_datetime_path)
shutil.copy2(script_path, latest_folder_path)

# Move the output datetime folder to the archive folder
shutil.move(output_datetime_path, archive_folder_path)
print(f"Output archived to: {archive_folder_path}")