import logging
from pathlib import Path

import numpy as np
import torch
from tqdm import tqdm

from train.train import (
    _baseline_anyview,
    _baseline_hard_vote,
    _forward_weights,
    _masked_mean_logits,
    _prepare_batch,
    _prepare_targets_cpu,
    evaluate,
)
from utils.compute_metrics import compute_metrics
from utils.save_predictions_to_las import save_predictions_to_las
from utils.utilities import compute_proba_batch


logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")


def test_deepchoice(model, test_loader, config, n_classes=11):
    return evaluate(model, test_loader, config, n_classes=n_classes, desc="Running Test")


def _reconstruct_coords(sample):
    coords_int = sample["coords_int"].cpu().numpy().astype(np.float64)
    coords_scale = sample["coords_scale"].cpu().numpy().astype(np.float64).reshape(-1, 1)
    coords_offset = sample["coords_offset"].cpu().numpy().astype(np.float64)
    return coords_int * coords_scale + coords_offset


def _normalized_weights(weights, mask):
    normalized = weights.masked_fill(~mask, float("-inf"))
    normalized = torch.softmax(normalized, dim=1)
    return torch.nan_to_num(normalized, nan=0.0)


def _metrics_dict(y_true, y_pred, n_classes, ignore_index):
    miou, mf1, ious = compute_metrics(y_true, y_pred, n_classes, ignore_index=ignore_index)
    return {"miou": float(miou), "mf1": float(mf1), "ious": ious}


def _predict_model_for_sample(model, sample, config):
    visibility, logits, mask, _ = _prepare_batch(sample, config)
    weights = _forward_weights(model, visibility, mask, config["model"]["type"])
    fused_logits = compute_proba_batch(weights, logits, mask=mask)
    pred = torch.argmax(fused_logits, dim=1)
    return pred, fused_logits, weights


def infer_deepchoice(model, test_loader, config, output_dir, comparison_model=None, comparison_config=None, comparison_field_name=None):
    model.eval()
    if comparison_model is not None:
        comparison_model.eval()
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    all_preds = []
    all_targets = []
    all_majority_preds = []
    all_hard_vote_preds = []
    all_anyview_preds = []
    all_comparison_preds = []
    tile_outputs = {}

    with torch.no_grad():
        for batch_idx, sample in enumerate(tqdm(test_loader, desc="Running Inference")):
            visibility, logits, mask, target = _prepare_batch(sample, config)
            weights = _forward_weights(model, visibility, mask, config["model"]["type"])
            fused_logits = compute_proba_batch(weights, logits, mask=mask)
            pred = torch.argmax(fused_logits, dim=1)
            majority_pred = torch.argmax(_masked_mean_logits(logits, mask), dim=1)
            hard_vote_pred = _baseline_hard_vote(logits, mask)
            anyview_pred = _baseline_anyview(logits, mask, target)
            norm_weights = _normalized_weights(weights, mask)
            comparison_pred = None
            if comparison_model is not None and comparison_config is not None:
                comparison_pred, _, _ = _predict_model_for_sample(comparison_model, sample, comparison_config)
            coords = _reconstruct_coords(sample)
            tile_names = sample["tile_name"]
            pred_np = pred.cpu().numpy()
            target_np = target.cpu().numpy()
            majority_np = majority_pred.cpu().numpy()
            hard_vote_np = hard_vote_pred.cpu().numpy()
            comparison_np = comparison_pred.cpu().numpy() if comparison_pred is not None else None
            unique_tiles = sorted(set(tile_names))
            for tile_name in unique_tiles:
                tile_mask = np.asarray([name == tile_name for name in tile_names], dtype=bool)
                if tile_name not in tile_outputs:
                    tile_outputs[tile_name] = {
                        "coords": [],
                        "pred": [],
                        "target": [],
                        "best_transformer": [],
                        "mean_prob_vote": [],
                        "hard_vote": [],
                    }
                    if comparison_field_name is not None:
                        tile_outputs[tile_name][comparison_field_name] = []
                tile_outputs[tile_name]["coords"].append(coords[tile_mask])
                tile_outputs[tile_name]["pred"].append(pred_np[tile_mask])
                tile_outputs[tile_name]["target"].append(target_np[tile_mask])
                tile_outputs[tile_name]["best_transformer"].append(pred_np[tile_mask])
                tile_outputs[tile_name]["mean_prob_vote"].append(majority_np[tile_mask])
                tile_outputs[tile_name]["hard_vote"].append(hard_vote_np[tile_mask])
                if comparison_field_name is not None and comparison_np is not None:
                    tile_outputs[tile_name][comparison_field_name].append(comparison_np[tile_mask])

            payload = {
                "pred": pred.cpu(),
                "target": target.cpu(),
                "fused_logits": fused_logits.cpu(),
                "weights": norm_weights.cpu(),
                "mask": sample["mask"].cpu(),
                "coords_int": sample["coords_int"].cpu(),
                "coords_scale": sample["coords_scale"].cpu(),
                "coords_offset": sample["coords_offset"].cpu(),
                "tile_name": sample["tile_name"],
                "source_path": sample["source_path"],
            }
            if "coords_tile_offset" in sample:
                payload["coords_tile_offset"] = sample["coords_tile_offset"].cpu()

            save_path = output_dir / f"pred_batch_{batch_idx:05d}.pt"
            torch.save(payload, save_path, pickle_protocol=4)

            all_preds.append(pred.cpu().numpy())
            all_targets.append(target.cpu().numpy())
            all_majority_preds.append(majority_pred.cpu().numpy())
            all_hard_vote_preds.append(hard_vote_pred.cpu().numpy())
            all_anyview_preds.append(anyview_pred.cpu().numpy())
            if comparison_pred is not None:
                all_comparison_preds.append(comparison_pred.cpu().numpy())

    y_true = np.concatenate(all_targets)
    y_pred = np.concatenate(all_preds)
    y_majority = np.concatenate(all_majority_preds)
    y_hard_vote = np.concatenate(all_hard_vote_preds)
    y_anyview = np.concatenate(all_anyview_preds)
    ignore_index = int(config["model"].get("ignore_index", 255))
    las_paths = []
    for tile_name, tile_payload in tile_outputs.items():
        coords = np.concatenate(tile_payload["coords"], axis=0)
        preds = np.concatenate(tile_payload["pred"], axis=0)
        targets = np.concatenate(tile_payload["target"], axis=0)
        extra_fields = {
            "best_transformer": np.concatenate(tile_payload["best_transformer"], axis=0),
            "mean_prob_vote": np.concatenate(tile_payload["mean_prob_vote"], axis=0),
            "hard_vote": np.concatenate(tile_payload["hard_vote"], axis=0),
        }
        if comparison_field_name is not None and comparison_field_name in tile_payload:
            extra_fields[comparison_field_name] = np.concatenate(tile_payload[comparison_field_name], axis=0)
        las_path = output_dir / f"{tile_name}_predictions.las"
        las_paths.append(
            save_predictions_to_las(
                coords,
                preds,
                las_path,
                ground_truth_array=targets,
                extra_fields=extra_fields,
            )
        )
    accuracy = float((y_true == y_pred).mean()) if y_true.size else float("nan")
    model_metrics = _metrics_dict(y_true, y_pred, int(config["model"]["num_classes"]), ignore_index)
    majority_metrics = _metrics_dict(y_true, y_majority, int(config["model"]["num_classes"]), ignore_index)
    hard_vote_metrics = _metrics_dict(y_true, y_hard_vote, int(config["model"]["num_classes"]), ignore_index)
    anyview_metrics = _metrics_dict(y_true, y_anyview, int(config["model"]["num_classes"]), ignore_index)
    logging.info(
        "Inference complete | saved %s batch predictions and %s LAS tiles | accuracy %.4f | model mIoU %.4f mF1 %.4f | majority %.4f %.4f | hard_vote %.4f %.4f | anyview %.4f %.4f",
        len(all_preds),
        len(las_paths),
        accuracy,
        model_metrics["miou"],
        model_metrics["mf1"],
        majority_metrics["miou"],
        majority_metrics["mf1"],
        hard_vote_metrics["miou"],
        hard_vote_metrics["mf1"],
        anyview_metrics["miou"],
        anyview_metrics["mf1"],
    )
    result = {
        "accuracy": accuracy,
        "num_batches": len(all_preds),
        "num_samples": int(y_true.size),
        "las_paths": las_paths,
        "metrics": {
            "model": model_metrics,
            "majority": majority_metrics,
            "hard_vote": hard_vote_metrics,
            "anyview": anyview_metrics,
        },
    }
    if all_comparison_preds and comparison_field_name is not None:
        y_comparison = np.concatenate(all_comparison_preds)
        result["metrics"][comparison_field_name] = _metrics_dict(
            y_true,
            y_comparison,
            int(config["model"]["num_classes"]),
            ignore_index,
        )
    return result


def test_with_baselines(model, test_loader, config, n_classes=11):
    metrics = test_deepchoice(model, test_loader, config, n_classes=n_classes)
    precomputed = config.get("test", {}).get("precomputed_baselines")
    if precomputed is not None:
        metrics["baselines"] = precomputed
        return metrics

    all_majority_preds = []
    all_hard_vote_preds = []
    all_anyview_preds = []
    all_targets = []
    max_views = int(config["model"]["max_views"])
    with torch.no_grad():
        for sample in tqdm(test_loader, desc="Running Test Baselines"):
            validate_target = _prepare_targets_cpu(sample["target"], config)
            logits = sample["logits"][:, :max_views, :]
            mask = sample["mask"][:, :max_views]
            majority_pred = torch.argmax(_masked_mean_logits(logits, mask), dim=1)
            hard_vote_pred = _baseline_hard_vote(logits, mask)
            anyview_pred = _baseline_anyview(logits, mask, validate_target)
            all_majority_preds.append(majority_pred.cpu().numpy())
            all_hard_vote_preds.append(hard_vote_pred.cpu().numpy())
            all_anyview_preds.append(anyview_pred.cpu().numpy())
            all_targets.append(validate_target.cpu().numpy())

    y_true = np.concatenate(all_targets)
    y_majority = np.concatenate(all_majority_preds)
    y_hard_vote = np.concatenate(all_hard_vote_preds)
    y_anyview = np.concatenate(all_anyview_preds)
    ignore_index = int(config["model"].get("ignore_index", 255))
    metrics["baselines"] = {
        "majority": _metrics_dict(y_true, y_majority, n_classes, ignore_index),
        "hard_vote": _metrics_dict(y_true, y_hard_vote, n_classes, ignore_index),
        "anyview": _metrics_dict(y_true, y_anyview, n_classes, ignore_index),
    }
    return metrics