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

import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.colors import LinearSegmentedColormap, ListedColormap
from PIL import Image

# --- Setup output folders ---
output_path = Path("/home/fedex/mt/plots/data_2d_projections_training")
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 as used in training (unstretched, grayscale)"
)
parser.add_argument(
    "--input1",
    type=Path,
    default=Path(
        "/home/fedex/mt/data/subter/new_projection/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/new_projection/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)",
)

args = parser.parse_args()

# --- Load the numpy projection data ---
proj_data1 = np.load(args.input1)
proj_data2 = np.load(args.input2)

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

# Debug info: Print the percentage of missing data in each frame
print(f"Frame 1 missing data percentage: {np.isnan(frame1).mean() * 100:.2f}%")
print(f"Frame 2 missing data percentage: {np.isnan(frame2).mean() * 100:.2f}%")

# --- Create a figure with 2 vertical subplots ---
fig, (ax1, ax2) = plt.subplots(nrows=2, ncols=1, figsize=(10, 5))

# Create custom colormap for missing data visualization
missing_color = [1, 0, 0, 1]  # Red with full alpha
cmap_missing = ListedColormap([missing_color])

# Replace the plotting section
for ax, frame, title in zip(
    (ax1, ax2),
    (frame1, frame2),
    (
        "Normal LiDAR Frame",
        "Degraded LiDAR Frame (Smoke)",
    ),
):
    # Create mask for missing data (directly from NaN values)
    missing_mask = np.isnan(frame)

    # Plot the valid data in grayscale
    frame_valid = np.copy(frame)
    frame_valid[missing_mask] = 0  # Set missing values to black in base image
    im = ax.imshow(frame_valid, cmap="gray", aspect="equal", vmin=0, vmax=0.8)

    # Overlay missing data in red with reduced alpha
    ax.imshow(
        missing_mask,
        cmap=ListedColormap([[1, 0, 0, 1]]),  # Pure red
        alpha=0.3,  # Reduced alpha for better visibility
    )

    ax.set_title(title)
    ax.axis("off")

# Adjust layout
plt.tight_layout()

# Create a more informative legend
legend_elements = [
    mpatches.Patch(facecolor="red", alpha=0.7, label="Missing Data"),
    mpatches.Patch(facecolor="white", label="Close Distance (0m)"),
    mpatches.Patch(facecolor="gray", label="Mid Distance"),
    mpatches.Patch(facecolor="black", label="Far Distance (70m)"),
]

# Add legend with better positioning and formatting
fig.legend(
    handles=legend_elements,
    loc="center right",
    bbox_to_anchor=(0.98, 0.5),
    title="Distance Information",
    framealpha=1.0,
)

# Save the plot
output_file = output_datetime_path / "data_2d_projections_training.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 training images ---
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}_training_{degradation_status}.tiff"
    )
    gray_image.save(gray_output_file)
    print(f"Training image saved to: {gray_output_file}")

# --- Handle folder structure ---
shutil.rmtree(latest_folder_path, ignore_errors=True)
latest_folder_path.mkdir(exist_ok=True, parents=True)

for file in output_datetime_path.iterdir():
    shutil.copy2(file, latest_folder_path)

script_path = Path(__file__)
shutil.copy2(script_path, output_datetime_path)
shutil.copy2(script_path, latest_folder_path)

shutil.move(output_datetime_path, archive_folder_path)
print(f"Output archived to: {archive_folder_path}")