import multiprocessing
import os
from pathlib import Path

import h5py
import numpy as np
import torch
from tqdm import tqdm


def _quantize_coords_batch(coords, coord_scale, coord_offset):
    coords = np.asarray(coords, dtype=np.float64)
    coord_offset = np.asarray(coord_offset, dtype=np.float64)
    return np.rint((coords - coord_offset) / coord_scale).astype(np.int32)


def _save_batch(save_path, vis, logits, masks, targets, coords, vmin, vmax, coord_scale, coord_offset, tile_offset=None):
    if tile_offset is None:
        tile_offset = np.zeros(3, dtype=np.float64)
    tile_offset = np.asarray(tile_offset, dtype=np.float64)
    coord_offset = np.asarray(coord_offset, dtype=np.float64)
    coords_global = np.asarray(coords, dtype=np.float64) + tile_offset
    coords_offset_global = coord_offset + tile_offset
    batch = {
        "visibility": torch.tensor(np.array(vis), dtype=torch.float32),
        "logits": torch.tensor(np.array(logits), dtype=torch.float32),
        "mask": torch.tensor(np.array(masks), dtype=torch.bool),
        "target": torch.tensor(np.array(targets), dtype=torch.long),
        "coords_int": torch.tensor(_quantize_coords_batch(coords_global, coord_scale, coords_offset_global), dtype=torch.int32),
        "coords_scale": torch.tensor(coord_scale, dtype=torch.float64),
        "coords_offset": torch.tensor(coords_offset_global, dtype=torch.float64),
        "coords_tile_offset": torch.tensor(tile_offset, dtype=torch.float64),
    }
    batch["visibility"] = normalize_visibility(batch["visibility"], vmin, vmax)
    save_path.parent.mkdir(parents=True, exist_ok=True)
    torch.save(batch, save_path, pickle_protocol=4)


def decode_coordinates(dataset, handle):
    coords = np.asarray(dataset)
    if np.issubdtype(coords.dtype, np.integer) and "coords_scale" in handle.attrs:
        scale = float(handle.attrs["coords_scale"])
        offset = np.asarray(handle.attrs.get("coords_offset", np.zeros(3, dtype=np.float32)), dtype=np.float32)
        return coords.astype(np.float32) * scale + offset
    return coords.astype(np.float32)


def decode_visibility(dataset, handle):
    visibility = np.asarray(dataset)
    if np.issubdtype(visibility.dtype, np.integer) and "visibility_quant_max" in handle.attrs:
        quant_max = float(handle.attrs["visibility_quant_max"])
        vmin = np.asarray(handle.attrs["visibility_vmin"], dtype=np.float32)
        vmax = np.asarray(handle.attrs["visibility_vmax"], dtype=np.float32)
        normalized = visibility.astype(np.float32) / max(quant_max, 1.0)
        return normalized * (vmax - vmin) + vmin
    return visibility.astype(np.float32)


def decode_logits(dataset, handle):
    logits = np.asarray(dataset)
    if np.issubdtype(logits.dtype, np.integer) and "logits_quant_max" in handle.attrs:
        quant_max = float(handle.attrs["logits_quant_max"])
        return logits.astype(np.float32) / max(quant_max, 1.0)
    return logits.astype(np.float32)


def decode_dense_visibility(visibility, observations):
    quant_max = float(observations["visibility_quant_max"])
    vmin = np.asarray(observations["visibility_vmin"], dtype=np.float32)
    vmax = np.asarray(observations["visibility_vmax"], dtype=np.float32)
    normalized = visibility.astype(np.float32) / max(quant_max, 1.0)
    return normalized * (vmax - vmin) + vmin


def decode_dense_logits(logits, observations):
    quant_max = float(observations["logits_quant_max"])
    return logits.astype(np.float32) / max(quant_max, 1.0)


def normalize_visibility(vis, vmin, vmax, eps=1e-8):
    mins = torch.tensor(vmin, dtype=torch.float32).view(1, 1, -1)
    maxs = torch.tensor(vmax, dtype=torch.float32).view(1, 1, -1)
    vis_clamped = torch.max(torch.min(vis, maxs), mins)
    return (vis_clamped - mins) / (maxs - mins + eps)


def _pad_views(array, max_views, pad_value=0, dtype=np.float32):
    current_views = array.shape[0]
    if current_views >= max_views:
        return array[:max_views]
    pad_shape = (max_views - current_views,) + array.shape[1:]
    padding = np.full(pad_shape, pad_value, dtype=dtype)
    return np.concatenate([array, padding], axis=0)


def process_h5_file(args):
    path, output_dir, vmin, vmax, batch_size, max_views, coord_scale, train, val, test = args
    batch_idx = 0
    filename_prefix = Path(path).stem

    with h5py.File(path, "r") as handle:
        points = handle["points"]
        keys = list(points.keys())
        for start in range(0, len(keys), batch_size):
            batch_keys = keys[start : start + batch_size]
            vis = []
            logits = []
            masks = []
            targets = []
            coords = []

            for key in batch_keys:
                point = points[key]
                visibility = decode_visibility(point["visibility"], handle)
                logit_vectors = decode_logits(point["logit_vectors"], handle)
                num_views = min(len(visibility), max_views)

                vis.append(_pad_views(visibility, max_views, pad_value=0.0, dtype=np.float32))
                logits.append(_pad_views(logit_vectors, max_views, pad_value=0.0, dtype=np.float32))

                mask = np.zeros(max_views, dtype=bool)
                mask[:num_views] = True
                masks.append(mask)

                targets.append(point["ground_truth"][()])
                coords.append(decode_coordinates(point["coordinates"], handle))

            if path in train:
                save_path = Path(output_dir) / "train" / f"{filename_prefix}_batch_{batch_idx:05d}.pt"
            elif path in test:
                save_path = Path(output_dir) / "test" / f"{filename_prefix}_batch_{batch_idx:05d}.pt"
            else:
                save_path = Path(output_dir) / "val" / f"{filename_prefix}_batch_{batch_idx:05d}.pt"

            coord_offset = np.min(np.asarray(coords, dtype=np.float32), axis=0)
            _save_batch(save_path, vis, logits, masks, targets, coords, vmin, vmax, coord_scale, coord_offset, np.zeros(3, dtype=np.float64))
            batch_idx += 1

    return f"{path} done"


def save_tile_observations_to_pt(
    cfg,
    tile_name,
    split_name,
    coords_array,
    tile_offset,
    all_observations,
    output_dir,
    batch_size=1024,
):
    output_dir = Path(output_dir) / split_name
    output_dir.mkdir(parents=True, exist_ok=True)
    max_views = int(cfg["selection"]["max_views"])
    min_views = int(cfg["selection"]["min_views"])
    vmin = cfg["data"]["vmin"]
    vmax = cfg["data"]["vmax"]
    batch_idx = 0
    coord_scale = float(cfg.get("storage", {}).get("coord_scale", 0.001))

    vis = []
    logits = []
    masks = []
    targets = []
    coords = []
    saved_paths = []

    if isinstance(all_observations, dict) and all_observations.get("mode") == "dense":
        counts = all_observations["counts"]
        ground_truth = all_observations["ground_truth"]
        for point_idx in range(len(counts)):
            num_views = int(counts[point_idx])
            if num_views < min_views:
                continue
            gt = int(ground_truth[point_idx])
            if gt < 0:
                continue

            visibility = decode_dense_visibility(all_observations["visibility"][point_idx, :num_views], all_observations)
            logit_vectors = decode_dense_logits(all_observations["logit_vectors"][point_idx, :num_views], all_observations)

            vis.append(_pad_views(visibility, max_views, pad_value=0.0, dtype=np.float32))
            logits.append(_pad_views(logit_vectors, max_views, pad_value=0.0, dtype=np.float32))

            mask = np.zeros(max_views, dtype=bool)
            mask[:num_views] = True
            masks.append(mask)
            targets.append(gt)
            coords.append(np.asarray(coords_array[point_idx], dtype=np.float32))

            if len(vis) == batch_size:
                save_path = output_dir / f"{tile_name}_batch_{batch_idx:05d}.pt"
                coord_offset = np.min(np.asarray(coords, dtype=np.float32), axis=0)
                _save_batch(save_path, vis, logits, masks, targets, coords, vmin, vmax, coord_scale, coord_offset, tile_offset)
                saved_paths.append(str(save_path))
                vis, logits, masks, targets, coords = [], [], [], [], []
                batch_idx += 1
    else:
        for point_idx, observations in all_observations.items():
            if len(observations["camera"]) < min_views:
                continue
            gt = observations["ground_truth"]
            if gt is None:
                continue

            visibility = np.asarray(observations["visibility"], dtype=np.float32)
            logit_vectors = np.asarray(observations["logit_vectors"], dtype=np.float32)
            num_views = min(len(visibility), max_views)

            vis.append(_pad_views(visibility, max_views, pad_value=0.0, dtype=np.float32))
            logits.append(_pad_views(logit_vectors, max_views, pad_value=0.0, dtype=np.float32))

            mask = np.zeros(max_views, dtype=bool)
            mask[:num_views] = True
            masks.append(mask)
            targets.append(gt)
            coords.append(np.asarray(coords_array[point_idx], dtype=np.float32))

            if len(vis) == batch_size:
                save_path = output_dir / f"{tile_name}_batch_{batch_idx:05d}.pt"
                coord_offset = np.min(np.asarray(coords, dtype=np.float32), axis=0)
                _save_batch(save_path, vis, logits, masks, targets, coords, vmin, vmax, coord_scale, coord_offset, tile_offset)
                saved_paths.append(str(save_path))
                vis, logits, masks, targets, coords = [], [], [], [], []
                batch_idx += 1

    if vis:
        save_path = output_dir / f"{tile_name}_batch_{batch_idx:05d}.pt"
        coord_offset = np.min(np.asarray(coords, dtype=np.float32), axis=0)
        _save_batch(save_path, vis, logits, masks, targets, coords, vmin, vmax, coord_scale, coord_offset, tile_offset)
        saved_paths.append(str(save_path))

    return saved_paths


def save_compact_payloads_to_pt(
    cfg,
    tile_name,
    split_name,
    coords_array,
    tile_offset,
    payloads,
    output_dir,
    batch_size=1024,
):
    output_dir = Path(output_dir) / split_name
    output_dir.mkdir(parents=True, exist_ok=True)
    max_views = int(cfg["selection"]["max_views"])
    min_views = int(cfg["selection"]["min_views"])
    vmin = cfg["data"]["vmin"]
    vmax = cfg["data"]["vmax"]

    payloads = [payload for payload in payloads if payload is not None and payload.get("point_indices") is not None]
    if not payloads:
        return []

    point_indices = np.concatenate([payload["point_indices"] for payload in payloads], axis=0)
    if point_indices.size == 0:
        return []
    visibility = np.concatenate([payload["visibility"] for payload in payloads], axis=0)
    logits = np.concatenate([payload["logit_vectors"] for payload in payloads], axis=0)
    ground_truth = np.concatenate([payload["ground_truth"] for payload in payloads], axis=0)

    order = np.argsort(point_indices, kind="mergesort")
    point_indices = point_indices[order]
    visibility = visibility[order]
    logits = logits[order]
    ground_truth = ground_truth[order]

    quant_max_vis = float(payloads[0]["visibility_quant_max"])
    quant_max_logits = float(payloads[0]["logits_quant_max"])
    payload_vmin = np.asarray(payloads[0]["visibility_vmin"], dtype=np.float32)
    payload_vmax = np.asarray(payloads[0]["visibility_vmax"], dtype=np.float32)

    saved_paths = []
    vis_batch = []
    logits_batch = []
    masks_batch = []
    targets_batch = []
    coords_batch = []
    batch_idx = 0
    coord_scale = float(cfg.get("storage", {}).get("coord_scale", 0.001))

    group_starts = np.r_[0, np.flatnonzero(np.diff(point_indices)) + 1]
    group_ends = np.r_[group_starts[1:], point_indices.size]

    for start, end in zip(group_starts, group_ends):
        point_idx = int(point_indices[start])
        num_obs = end - start
        if num_obs < min_views:
            continue

        gt_slice = ground_truth[start:end]
        valid_gt = gt_slice[gt_slice >= 0]
        if valid_gt.size == 0:
            continue
        gt = int(valid_gt[0])

        local_visibility = visibility[start:end]
        if num_obs > max_views:
            keep = np.lexsort((local_visibility[:, 0], local_visibility[:, 1]))[:max_views]
            local_visibility = local_visibility[keep]
            local_logits = logits[start:end][keep]
            num_views = max_views
        else:
            local_logits = logits[start:end]
            num_views = num_obs

        vis_float = (local_visibility.astype(np.float32) / max(quant_max_vis, 1.0)) * (payload_vmax - payload_vmin) + payload_vmin
        logits_float = local_logits.astype(np.float32) / max(quant_max_logits, 1.0)

        vis_batch.append(_pad_views(vis_float, max_views, pad_value=0.0, dtype=np.float32))
        logits_batch.append(_pad_views(logits_float, max_views, pad_value=0.0, dtype=np.float32))
        mask = np.zeros(max_views, dtype=bool)
        mask[:num_views] = True
        masks_batch.append(mask)
        targets_batch.append(gt)
        coords_batch.append(np.asarray(coords_array[point_idx], dtype=np.float32))

        if len(vis_batch) == batch_size:
            save_path = output_dir / f"{tile_name}_batch_{batch_idx:05d}.pt"
            coord_offset = np.min(np.asarray(coords_batch, dtype=np.float32), axis=0)
            _save_batch(
                save_path,
                vis_batch,
                logits_batch,
                masks_batch,
                targets_batch,
                coords_batch,
                vmin,
                vmax,
                coord_scale,
                coord_offset,
                tile_offset,
            )
            saved_paths.append(str(save_path))
            vis_batch, logits_batch, masks_batch, targets_batch, coords_batch = [], [], [], [], []
            batch_idx += 1

    if vis_batch:
        save_path = output_dir / f"{tile_name}_batch_{batch_idx:05d}.pt"
        coord_offset = np.min(np.asarray(coords_batch, dtype=np.float32), axis=0)
        _save_batch(
            save_path,
            vis_batch,
            logits_batch,
            masks_batch,
            targets_batch,
            coords_batch,
            vmin,
            vmax,
            coord_scale,
            coord_offset,
            tile_offset,
        )
        saved_paths.append(str(save_path))

    return saved_paths


def save_dataset_to_pt_parallel(cfg, train, val, test, output_dir, batch_size=1024, num_workers=4):
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    max_views = cfg["selection"]["max_views"]
    vmin = cfg["data"]["vmin"]
    vmax = cfg["data"]["vmax"]
    h5_paths = train + val + test
    args_list = [
        (path, str(output_dir), vmin, vmax, batch_size, max_views, float(cfg.get("storage", {}).get("coord_scale", 0.001)), train, val, test)
        for path in h5_paths
    ]

    with multiprocessing.Pool(num_workers) as pool:
        for result in pool.imap_unordered(process_h5_file, args_list):
            print(result)