import os
import time
import warnings
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from typing import Dict, List

import numpy as np
import torch
import yaml
# uv pip install lpips scikit-image
from lpips import LPIPS
from skimage.metrics import structural_similarity

from anydata.geometry.depth import DEPTH18_SCALE, DEPTH24_SCALE, decode_depth, encode_depth
from anydata.utils.read import read_image_seq, read_npz_seq, read_npzs_seq, read_video_seq, read_zarr_seq
from anydata.utils.write import write_image_seq, write_npz, write_npzs_seq, write_video_seq, write_zarr_seq

warnings.filterwarnings("ignore", category=UserWarning)


def create_moving_rainbow_array(frames=60, h=256, w=256):
    # 1. Create a normalized diagonal coordinate grid (0.0 to 1.0)
    x = np.linspace(0, 1, w)
    y = np.linspace(0, 1, h)
    xv, yv = np.meshgrid(x, y)
    
    # Diagonal position: (x + y) / 2 creates a 45-degree gradient
    diagonal = (xv + yv) / 2
    
    video_data = []
    for i in range(frames):
        # 2. Shift the hue over time
        shift = i / frames
        hue = (diagonal + shift) % 1.0
        
        # 3. Fast HSV to RGB conversion for S=1, V=1
        # This creates the classic ROYGBIV rainbow spectrum
        r = np.clip(np.abs(hue * 6 - 3) - 1, 0, 1)
        g = np.clip(2 - np.abs(hue * 6 - 2), 0, 1)
        b = np.clip(2 - np.abs(hue * 6 - 4), 0, 1)
        
        # Stack into (H, W, 3) and convert to 8-bit
        rgb_frame = (np.stack([r, g, b], axis=-1) * 255).astype(np.uint8)
        video_data.append(rgb_frame)

    return np.stack(video_data)


@contextmanager
def measure_time(output_dict: dict, label: str):
    start_time = time.perf_counter()
    yield
    output_dict[label] = time.perf_counter() - start_time


def get_size(path):
    # Base Case: Path is a single file
    if os.path.isfile(path):
        return os.path.getsize(path) / 1e6  # Convert bytes to megabytes
    # Recursive Case: Path is a directory
    elif os.path.isdir(path) and not os.path.islink(path):
        return sum(get_size(os.path.join(path, entry)) for entry in os.listdir(path))
    return 0 # Path does not exist or is not a file/folder


def format_floats_recursive(data, sigfigs=4):
    if isinstance(data, float):
        return float(np.round(data, decimals=2)) if data>100 else float(format(data, f".{sigfigs}g"))
    elif isinstance(data, dict):
        return {k: format_floats_recursive(v, sigfigs) for k, v in data.items()}
    else:
        return data


def lpips(original: np.ndarray, compressed: np.ndarray, device: str = "cpu", batch_size: int = None):
    """Lower is better. 0 = identical"""
    device = torch.device(device if torch.cuda.is_available() else 'cpu')
    images = np.stack([original, compressed])
    image_tensor = torch.from_numpy(images).float() / 127.5 - 1
    if image_tensor.ndim == 4:  # [2, H, W, 3] -> [2, 1, H, W, 3]
        image_tensor = image_tensor.unsqueeze(1)
    image_tensor = image_tensor.permute(0, 1, 4, 2, 3)  # (2, N, C, H, W)
    
    batch_size = min(batch_size, image_tensor.shape[1]) if batch_size else image_tensor.shape[1]
    loss_fn = LPIPS(net='alex', verbose=False).to(device)  # or 'vgg', 'squeeze'
    with torch.no_grad():
        distance = []
        for i in range(0, image_tensor.shape[1], batch_size):
            batch = image_tensor[:, i:i+batch_size].to(device)
            distance.append(loss_fn(batch[0], batch[1]))
        distance = torch.cat(distance)
    return distance.mean().item()


def ssim(original: np.ndarray, compressed: np.ndarray):
    """Higher is better. Infinity = identical"""
    images = np.stack([original, compressed])
    if images.ndim == 4:  # [2, H, W, 3] -> [2, 1, H, W, 3]
        images = images[:, np.newaxis]

    def ssim_frame(i):
        return structural_similarity(images[0, i], images[1, i], channel_axis=2, data_range=255)
    
    with ThreadPoolExecutor() as executor:
        scores = list(executor.map(ssim_frame, range(original.shape[0])))
    return np.mean(scores).item()


def psnr(original: np.ndarray, compressed: np.ndarray, data_range=255.0):
    """
    Computes PSNR for a batch of images.
    img_true, img_test: NumPy arrays of shape (N, H, W, C)
    """
    images = np.stack([original, compressed]).astype(np.float32)
    if images.ndim == 4:  # [2, H, W, 3] -> [2, 1, H, W, 3]
        images = images[:, np.newaxis]
    mse = np.mean((images[0] - images[1])**2, axis=(1, 2, 3))
    psnr = np.full_like(mse, 100, dtype=np.float32)
    nonzero_mse = mse > 0
    psnr[nonzero_mse] = 10 * np.log10(data_range**2 / mse[nonzero_mse])
    return np.mean(psnr).item()


def compute_image_metrics(original_rgb: np.ndarray, test_rgb: np.ndarray):
    results = {}
    results["pixel"] = np.mean(np.abs(original_rgb - test_rgb)).item()
    results["LPIPS"] = lpips(original_rgb, test_rgb, "cuda", 20)
    results["SSIM"] = ssim(original_rgb, test_rgb)
    results["PSNR"] = psnr(original_rgb, test_rgb)
    return results


def test_rgb(filename: str, rgb_array: np.ndarray, write_fn, read_fn, **kwargs):
    results = {"time": {}, "similarity": {}}

    with measure_time(results["time"], "write"):
        path = write_fn(filename, rgb_array, **kwargs)
    
    with measure_time(results["time"], "read_full"):
        read_data_full: np.ndarray = read_fn(filename)

    with measure_time(results["time"], "read_seq"):
        read_frames = [read_fn(filename, index=i) for i in range(len(rgb_array))]
    read_frames_stacked = np.stack(read_frames)
    
    assert np.array_equal(read_frames_stacked, read_data_full), "Frame-by-frame read does not match full read"
    
    results["similarity"] = compute_image_metrics(rgb_array, read_data_full)
    results["size"] = get_size(path)
    return results


def test_lowdim(filename: str, lowdim: Dict[str, np.ndarray], write_fn, read_fn):
    results = {"time": {}}
    seq_len = next(iter(lowdim.values())).shape[0]

    with measure_time(results["time"], "write"):
        path = write_fn(filename, lowdim)
    
    with measure_time(results["time"], "read_full"):
        read_data_full: np.ndarray = read_fn(filename)

    with measure_time(results["time"], "read_seq"):
        read_frames: List[dict] = [read_fn(filename, index=i) for i in range(seq_len)]
    stacked_values = list(map(np.stack, zip(*(f.values() for f in read_frames))))
    read_frames_stacked = dict(zip(read_frames[0].keys(), stacked_values))
    
    for key, val in lowdim.items():
        assert np.array_equal(read_data_full[key], val)
        assert np.array_equal(read_frames_stacked[key], val)
    
    results["size"] = get_size(path)
    return results


def test_depth(filename: str, depth: np.ndarray, write_fn, read_fn, nbits: int = 32, tol: float = 0):
    results = {"time": {}}

    with measure_time(results["time"], "write"):
        path = write_fn(filename, {"depth": encode_depth(depth, bits=nbits)})
    
    with measure_time(results["time"], "read_full"):
        read_data_full: np.ndarray = decode_depth(read_fn(filename, key="depth"), bits=nbits)

    with measure_time(results["time"], "read_seq"):
        read_frames = [decode_depth(read_fn(filename, key="depth", index=i), bits=nbits) for i in range(len(depth))]
    read_frames_stacked = np.stack(read_frames)
    
    assert np.allclose(read_data_full, depth, atol=tol)
    assert np.allclose(read_frames_stacked, depth, atol=tol)
    
    results["error"] = float(np.max(np.abs(read_data_full - depth)))
    results["size"] = get_size(path)
    return results


def test_nested(filename: str, nested_dict: dict, write_fn, read_fn, test_key: str = None):
    results = {"time": {}}

    with measure_time(results["time"], "write"):
        path = write_fn(filename, nested_dict)
    
    with measure_time(results["time"], "read_full"):
        read_data_full: dict = read_fn(filename)

    def compare_dicts(d1, d2, slice1=slice(None), slice2=slice(None)):
        for key in d1:
            assert key in d2, f"Key '{key}' missing in read data"
            if isinstance(d1[key], dict):
                compare_dicts(d1[key], d2[key], slice1, slice2)
            else:
                assert np.array_equal(d1[key][slice1], d2[key][slice2]), f"Mismatch at key '{key}'"

    compare_dicts(nested_dict, read_data_full)

    read_slice = read_fn(filename, key=test_key)
    test_value = nested_dict
    for key in test_key.split('/'):
        test_value = test_value[key]
    compare_dicts(test_value, read_slice)
    
    index_slice = slice(10, 20, 2)
    read_slice = read_fn(filename, index=index_slice)
    compare_dicts(nested_dict, read_slice, slice1=index_slice, slice2=slice(None))

    index_list = [20, 40]
    read_list = read_fn(filename, index=index_list)
    compare_dicts(nested_dict, read_list, slice1=index_list, slice2=slice(None))

    single_data = read_fn(filename, index=5)
    compare_dicts(nested_dict, single_data, slice1=5, slice2=())
    
    results["size"] = get_size(path)
    return results


def main():
    num_steps, height, width = 50, 1082, 1920
    rgb_data = create_moving_rainbow_array(frames=num_steps, h=height, w=width)
    depth_data = ((2**18-1) / DEPTH18_SCALE) * np.random.rand(num_steps, height, width).astype(np.float32)
    lowdim_data = {
        "intrinsics": np.random.rand(num_steps, 3, 3).astype(np.float32),
        "extrinsics": np.random.rand(num_steps, 4, 4).astype(np.float32),
        "timestep": np.arange(num_steps).astype(np.int64),
    }
    nested_dict_data = {
        "level1": {
            "level2": {
                "intrinsics": np.random.rand(num_steps, 3, 3).astype(np.float32),
                "extrinsics": np.random.rand(num_steps, 4, 4).astype(np.float32),
            },
            "timestep": np.arange(num_steps).astype(np.int64),
        },
        "test": np.random.rand(num_steps, 6, 8).astype(np.float32)
    }

    dst = "debug/target"
    os.makedirs(dst, exist_ok=True)
    test_nested(f"{dst}/nested", nested_dict_data, write_npz, read_npz_seq, "level1/level2")
    test_nested(f"{dst}/nested", nested_dict_data, write_npzs_seq, read_npzs_seq, "level1/level2")
    test_nested(f"{dst}/nested", nested_dict_data, write_zarr_seq, read_zarr_seq, "level1/level2")

    print("Testing low-dimensional data compression...")
    ld_npz = test_lowdim(f"{dst}/lowdim", lowdim_data, write_npz, read_npz_seq)
    ld_npzs = test_lowdim(f"{dst}/lowdim", lowdim_data, write_npzs_seq, read_npzs_seq)
    ld_zarr = test_lowdim(f"{dst}/lowdim", lowdim_data, write_zarr_seq, read_zarr_seq)

    print("Testing RGB compression...")
    rgb_jpg = test_rgb(f"{dst}/rgb", rgb_data, write_image_seq, read_image_seq)
    rgb_mp4 = test_rgb(f"{dst}/rgb", rgb_data, write_video_seq, read_video_seq)
    rgb_mp4_444 = test_rgb(f"{dst}/rgb_444", rgb_data, write_video_seq, read_video_seq, pixel_format='yuv444p')
    
    print("Testing depth data compression...")
    depth_npz = test_depth(f"{dst}/depth", depth_data, write_npz, read_npz_seq)
    depth_npzs = test_depth(f"{dst}/depth", depth_data, write_npzs_seq, read_npzs_seq)
    depth_zarr = test_depth(f"{dst}/depth", depth_data, write_zarr_seq, read_zarr_seq)
    depth24_npzs = test_depth(f"{dst}/depth24", depth_data, write_npzs_seq, read_npzs_seq, 24, 1/DEPTH24_SCALE)
    depth24_zarr = test_depth(f"{dst}/depth24", depth_data, write_zarr_seq, read_zarr_seq, 24, 1/DEPTH24_SCALE)
    depth18_npzs = test_depth(f"{dst}/depth18", depth_data, write_npzs_seq, read_npzs_seq, 18, 1/DEPTH18_SCALE)
    depth18_zarr = test_depth(f"{dst}/depth18", depth_data, write_zarr_seq, read_zarr_seq, 18, 1/DEPTH18_SCALE)
    
    all_results = format_floats_recursive({
        "rgb": {"jpg": rgb_jpg, "mp4": rgb_mp4, "mp4_444": rgb_mp4_444},
        "lowdim": {"npzs": ld_npzs, "npz": ld_npz, "zarr": ld_zarr},
        "depth": {
            "32-bit": {"npzs": depth_npzs, "npz": depth_npz, "zarr": depth_zarr},
            "24-bit": {"npzs": depth24_npzs, "zarr": depth24_zarr},
            "18-bit": {"npzs": depth18_npzs, "zarr": depth18_zarr},
        }
    })

    with open(f"{dst}/results.yaml", "w") as f:
        yaml.dump(all_results, f, sort_keys=False, default_flow_style=False)
    print(f"\nTesting done. Compression Results written to {dst}/results.yaml:")


if __name__ == "__main__":
    main()