import json
import math
import shutil
from collections import defaultdict
from pathlib import Path
from decimal import Decimal, ROUND_HALF_UP

import colorful
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
import torch
import trimesh
from scipy.spatial import cKDTree as KDTree
from tabulate import tabulate

from pcn import click, load_model, mkdir
from pcn import point_cloud as pcl
from pcn.datasets.uniform_sparse_sampling import (Dataset, MeshDataDescription,
                                                  MinibatchGenerator,
                                                  PointCloudAndMeshPairDataset,
                                                  GtUniformPointCloudData)
from pcn.experiment import (Model, ModelHyperparameters,
                            TrainingHyperparameters, setup_model)


def _compute_chamfer_distance(gt_surface_points: np.ndarray,
                              pred_surface_points: np.ndarray):
    gt_points_kd_tree = KDTree(gt_surface_points)
    distances, locations = gt_points_kd_tree.query(pred_surface_points)
    cd_term1 = np.mean(np.square(distances))

    pred_points_kd_tree = KDTree(pred_surface_points)
    distances, locations = pred_points_kd_tree.query(gt_surface_points)
    cd_term2 = np.mean(np.square(distances))

    return cd_term1 + cd_term2


def _compute_non_squared_chamfer_distance(gt_surface_points: np.ndarray,
                                          pred_surface_points: np.ndarray):
    gt_points_kd_tree = KDTree(gt_surface_points)
    distances, locations = gt_points_kd_tree.query(pred_surface_points)
    cd_term1 = np.mean(distances)

    pred_points_kd_tree = KDTree(pred_surface_points)
    distances, locations = pred_points_kd_tree.query(gt_surface_points)
    cd_term2 = np.mean(distances)

    return cd_term1 + cd_term2


def _summarize(result_directory: str):
    result_directory = Path(result_directory)
    metrics_directory = result_directory / "metrics"
    args_path = result_directory / "args.json"
    assert args_path.is_file()
    with open(args_path) as f:
        args = json.load(f)
    table = defaultdict(list)
    metrics_path_list = list(metrics_directory.glob("*.json"))
    for metrics_path in metrics_path_list:
        try:
            with open(metrics_path) as f:
                metrics = json.load(f)
        except json.decoder.JSONDecodeError:
            print("Error:", metrics_path)
            continue
        for key, value in metrics.items():
            table[key].append(value)

    return table, args, len(metrics_path_list)


def _summarize(result_directory: str):
    result_directory = Path(result_directory)
    metrics_directory = result_directory / "metrics"
    table = defaultdict(list)
    metrics_path_list = list(metrics_directory.glob("*.json"))
    for metrics_path in metrics_path_list:
        try:
            with open(metrics_path) as f:
                metrics = json.load(f)
        except json.decoder.JSONDecodeError:
            print("Error:", metrics_path)
            continue
        for key, value in metrics.items():
            table[key].append(value)

    return table, len(metrics_path_list)


@click.group()
def client():
    pass


@client.command()
@click.argument("--result-directory", type=str, required=True)
def summarize(args):
    metric_list = ["chamfer_distance"]
    settings_list = ["num_input_points"]
    result, num_data = _summarize(args.result_directory)
    print(colorful.bold("Hyperparameters:"))
    print(args)

    data = {
        "num_data": num_data,
    }
    for metric_name in metric_list:
        values = [
            value for value in result[metric_name]
            if value is not None and value > 0
        ]
        values = np.array(values)
        ignore = values == np.nan
        print(values.dtype)
        print(np.count_nonzero(ignore), "nans")
        values = values[~ignore]
        print(values.shape, values.dtype)
        print(np.min(values), np.max(values))
        print(colorful.bold("Histogram:"))
        print(np.histogram(values))
        mean = np.mean(values)
        std = np.std(values)
        data[f"{metric_name}_mean"] = mean
        data[f"{metric_name}_std"] = std
    for key in settings_list:
        values = np.array(result[key])
        data[key] = values[0]

    tabulate_row = [data["num_data"]]
    for key in settings_list:
        value = data[key]
        tabulate_row.append(f"{value}")
    for metric_name in metric_list:
        mean = data[f"{metric_name}_mean"]
        std = data[f"{metric_name}_std"]
        mean = Decimal(mean).quantize(Decimal("0.001"), ROUND_HALF_UP)
        std = Decimal(std).quantize(Decimal("0.001"), ROUND_HALF_UP)
        tabulate_row.append(f"{mean:.03f} (±{std:.03f})")

    print(colorful.bold("Result:"))
    print(
        tabulate([tabulate_row],
                 headers=["# of data"] + settings_list + metric_list,
                 tablefmt="github"))


@client.command(name="chamfer_distance_surface_data")
@click.argument("--output-directory", type=str, required=True)
@click.argument("--npz-path", type=str, required=True)
@click.argument("--num-input-points", type=int, default=50)
@click.argument("--variance", type=float, default=0.02)
@click.argument("--chamfer-distance-method",
                type=click.Choice(
                    ["asymmetric", "symmetric", "non_squared_symmetric"]),
                required=True)
def chamfer_distance_surface_data(args):
    device = torch.device("cuda", 0)
    npz_path = Path(args.npz_path)
    output_directory = Path(args.output_directory)
    metrics_directory = output_directory / "metrics"
    mkdir(output_directory)
    mkdir(metrics_directory)

    path_list = [npz_path]
    dataset = Dataset(path_list)
    minibatch_generator = MinibatchGenerator(
        num_input_points=args.num_input_points, device=device)

    if args.chamfer_distance_method == "symmetric":
        compute_chamfer_distance = _compute_chamfer_distance
    elif args.chamfer_distance_method == "non_squared_symmetric":
        compute_chamfer_distance = _compute_non_squared_chamfer_distance
    else:
        raise NotImplementedError()

    for pc_data in dataset:
        batch = minibatch_generator([pc_data])
        num_samples = 16384
        indices = np.random.choice(args.num_input_points, size=num_samples)
        pred_dense_points = batch.input_points[0][indices].cpu().numpy()
        pred_dense_points += np.random.normal(scale=math.sqrt(args.variance),
                                              size=pred_dense_points.shape)
        gt_surface_points = batch.gt_dense_points.cpu().numpy()[0]

        gt_surface_points = gt_surface_points / pc_data.scale - pc_data.offset
        pred_dense_points = pred_dense_points / pc_data.scale - pc_data.offset

        print(batch.input_points[0].shape, "->", pred_dense_points.shape)
        print(gt_surface_points.shape, "<->", pred_dense_points.shape)

        chamfer_distance = compute_chamfer_distance(gt_surface_points,
                                                    pred_dense_points)

        parts = str(pc_data.path).split("/")
        category_id = parts[-3]
        model_id = parts[-2]
        print(category_id, model_id, chamfer_distance, flush=True)

        result_path = metrics_directory / f"{category_id}_{model_id}.json"
        print(chamfer_distance)
        print(result_path)

        result = {
            "model_id": f"{category_id}_{model_id}",
            "chamfer_distance": chamfer_distance,
            "num_input_points": args.num_input_points,
            "chamfer_distance_num_samples": num_samples,
            "chamfer_distance_method": args.chamfer_distance_method,
        }
        with open(result_path, "w") as f:
            json.dump(result, f, indent=4, sort_keys=True)


if __name__ == "__main__":
    client()
