import json
import math
import shutil
from collections import defaultdict
from pathlib import Path
from hashlib import sha256

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

from implicit_geometric_regularization import click, load_model, mkdir
from implicit_geometric_regularization import point_cloud as pcl
from implicit_geometric_regularization.datasets.shapenet.partial_sampling import (
    MinibatchGenerator, PartialPointCloudDataset, PartialPointCloudData)
from implicit_geometric_regularization.datasets.shapenet.uniform_sparse_sampling import \
    Minibatch
from implicit_geometric_regularization.experiments.learning_shape_space import (
    AdaptiveLatentOptimization, DecoderHyperparameters, LatentOptimization,
    LatentOptimizationHyperparameters, LossFunction, MarchingCubes, Model,
    ModelHyperparameters, TrainingHyperparameters, setup_model)


def _sample_surface_points(faces: np.ndarray, vertices: np.ndarray,
                           num_point_samples: int):
    mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
    samples, face_index = trimesh.sample.sample_surface(
        mesh, num_point_samples)
    return samples


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 _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:
        with open(metrics_path) as f:
            metrics = json.load(f)
        for key, value in metrics.items():
            table[key].append(value)

    return table, args, len(metrics_path_list)


class SurfaceChamferDistance:
    def __init__(self,
                 minibatch_generator: MinibatchGenerator,
                 optimize_latent_func: LatentOptimization,
                 model: Model,
                 grid_size: int,
                 latent_optimization_num_samples: int,
                 num_viewpoint_samples: int,
                 chamfer_distance_num_samples: int,
                 seed: int = 0):
        self.minibatch_generator = minibatch_generator
        self.optimize_latent = optimize_latent_func
        self.model = model
        self.grid_size = grid_size
        self.latent_optimization_num_samples = latent_optimization_num_samples
        self.chamfer_distance_num_samples = chamfer_distance_num_samples
        self.num_viewpoint_samples = num_viewpoint_samples
        self.seed = seed

    def __call__(self, pc_data: PartialPointCloudData):
        if self.seed == -1:
            random_state = np.random.RandomState(None)
        else:
            torch.manual_seed(self.seed)
            random_state = np.random.RandomState(self.seed)

        data = self.minibatch_generator(
            [pc_data],
            num_point_samples=self.latent_optimization_num_samples,
            num_viewpoint_samples=pc_data.num_viewpoints,
            random_state=random_state)

        device = self.model.get_device()
        gt_surface_points = pc_data.vertices / pc_data.scale - pc_data.offset

        if self.seed == -1:
            seed = pc_data.path.parts[-3] + "/" + pc_data.path.parts[-2]
            print(seed)
            hash = sha256(seed.encode())
            seed = np.frombuffer(hash.digest(), dtype=np.uint32)
            random_state = np.random.RandomState(seed)
        else:
            torch.manual_seed(self.seed)
            random_state = np.random.RandomState(self.seed)
        view_indices = random_state.choice(len(data.points_list),
                                           size=self.num_viewpoint_samples,
                                           replace=False)
        print(view_indices)
        chamfer_distance_for_each_viewpoints = []
        for view_index in view_indices:
            context_points = data.points_list[view_index]
            initial_z = torch.normal(mean=0,
                                     std=0.01,
                                     size=(1, self.model.z_map.shape[1]),
                                     dtype=torch.float32).to(device)

            input_data = Minibatch(points=context_points, data_index=None)
            z = self.optimize_latent(input_data, initial_z)

            grid_max_value = 1
            grid_min_value = -1
            marching_cubes = MarchingCubes(model=self.model,
                                           grid_size=self.grid_size,
                                           grid_max_value=grid_max_value,
                                           grid_min_value=grid_min_value)
            try:
                mc_vertices, mc_faces = marching_cubes(z)
                mc_vertices = mc_vertices / pc_data.scale - pc_data.offset
            except ValueError:
                chamfer_distance = None

            pred_surface_points = _sample_surface_points(
                mc_faces, mc_vertices, self.chamfer_distance_num_samples)

            chamfer_distance = _compute_chamfer_distance(
                gt_surface_points, pred_surface_points)

            chamfer_distance_for_each_viewpoints.append({
                "view_index":
                int(view_index),
                "chamfer_distance":
                chamfer_distance
            })
            print("view:",
                  view_index,
                  "chamfer_distance:",
                  chamfer_distance,
                  flush=True)

        return chamfer_distance_for_each_viewpoints


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


@client.command()
@click.argument("--result-directory", type=str, required=True)
def summarize(args):
    result, args, num_data = _summarize(args.result_directory)

    latent_optimization_num_samples = result[
        "latent_optimization_num_samples"][0]

    data_list = []
    for model_id, chamfer_distance_for_each_viewpoints in zip(
            result["model_id"],
            result["chamfer_distance_for_each_viewpoints"]):
        for row in chamfer_distance_for_each_viewpoints:
            view_index = row["view_index"]
            chamfer_distance = row["chamfer_distance"]
            if chamfer_distance is None:
                print(model_id, view_index, "failed")
                continue
            data_list.append({
                "model_id": model_id,
                "view_index": view_index,
                "chamfer_distance": chamfer_distance
            })

    df = pd.DataFrame(data_list)
    grouped = df.groupby("view_index")
    print(grouped)
    print(grouped.size())

    mean = df.mean(numeric_only=True)["chamfer_distance"]
    std = df.std(numeric_only=True)["chamfer_distance"]
    print(
        tabulate(
            [[latent_optimization_num_samples, f"{mean:.06f} (±{std:.06f})"]],
            headers=["# of context", "chamfer_distance"],
            tablefmt="github"))


@client.command(name="chamfer_distance_surface_data")
@click.argument("--checkpoint-directory", type=str, required=True)
@click.argument("--checkpoint-epoch", type=int, default=None)
@click.argument("--output-directory", type=str, required=True)
@click.argument("--pc-data-path", type=str, required=True)
@click.argument("--grid-size", type=int, default=256)
@click.argument("--num-input-points", type=int, default=50)
@click.argument("--num-viewpoint-samples", type=int, default=5)
@click.argument("--seed", type=int, default=0)
@click.argument("--chamfer-distance-num-samples", type=int, default=30000)
@click.argument("--skip-if-exists", is_flag=True)
@click.hyperparameter_class(LatentOptimizationHyperparameters)
def chamfer_distance_surface_data(
        args, latent_optimization_params: LatentOptimizationHyperparameters):
    device = torch.device("cuda", 0)
    pc_data_path = Path(args.pc_data_path)
    output_directory = Path(args.output_directory)
    metrics_directory = output_directory / "metrics"
    mkdir(output_directory)
    mkdir(metrics_directory)

    parts = pc_data_path.parts
    category_id = parts[-3]
    model_id = parts[-2]
    result_path = metrics_directory / f"{category_id}_{model_id}.json"
    if args.skip_if_exists and result_path.exists():
        print("skip", pc_data_path)
        return

    checkpoint_directory = Path(args.checkpoint_directory)
    args_path = checkpoint_directory / "args.json"
    if args.checkpoint_epoch is None:
        model_path = checkpoint_directory / "model.pt"
    else:
        model_path = checkpoint_directory / f"model.{args.checkpoint_epoch}.pt"
    assert args_path.is_file()
    assert model_path.is_file()

    shutil.copyfile(args_path, output_directory / "args.json")

    model_hyperparams = ModelHyperparameters.load_json(args_path)
    decoder_hyperparams = DecoderHyperparameters.load_json(args_path)
    training_hyperparams = TrainingHyperparameters.load_json(args_path)

    model = setup_model(model_hyperparams, decoder_hyperparams)
    load_model(model_path, model)
    model.to(device)
    model.eval()

    dataset = PartialPointCloudDataset([pc_data_path])
    minibatch_generator = MinibatchGenerator(device=device)

    loss_function = LossFunction(
        tau=training_hyperparams.loss_tau,
        lam=training_hyperparams.loss_lambda,
        alpha=training_hyperparams.loss_alpha,
        num_eikonal_samples=training_hyperparams.num_eikonal_samples,
        eikonal_term_default_stddev=training_hyperparams.eikonal_term_stddev)

    optimize_latent = AdaptiveLatentOptimization(
        model=model,
        loss_function=loss_function,
        params=latent_optimization_params)
    compute_chamfer_distance = SurfaceChamferDistance(
        minibatch_generator=minibatch_generator,
        optimize_latent_func=optimize_latent,
        model=model,
        num_viewpoint_samples=args.num_viewpoint_samples,
        grid_size=args.grid_size,
        latent_optimization_num_samples=args.num_input_points,
        chamfer_distance_num_samples=args.chamfer_distance_num_samples,
        seed=args.seed)

    for pc_data in dataset:
        chamfer_distance_for_each_viewpoints = compute_chamfer_distance(
            pc_data)

        parts = str(pc_data.path).split("/")
        category_id = parts[-3]
        model_id = parts[-2]
        result_path = metrics_directory / f"{category_id}_{model_id}.json"

        print(category_id, model_id, "done", flush=True)
        print(result_path)

        result = {
            "model_id": f"{category_id}_{model_id}",
            "chamfer_distance_for_each_viewpoints":
            chamfer_distance_for_each_viewpoints,
            "chamfer_distance_num_samples": args.chamfer_distance_num_samples,
            "grid_size": args.grid_size,
            "latent_optimization_num_samples":
            latent_optimization_params.num_samples,
            "latent_optimization_iterations":
            latent_optimization_params.iterations,
            "latent_optimization_initial_lr":
            latent_optimization_params.initial_lr,
            "latent_optimization_decrease_lr_every":
            latent_optimization_params.decrease_lr_every,
            "latent_optimization_lr_decay_factor":
            latent_optimization_params.lr_decay_factor,
            "num_viewpoint_samples": args.num_viewpoint_samples,
        }
        with open(result_path, "w") as f:
            json.dump(result, f, indent=4, sort_keys=True)


if __name__ == "__main__":
    client()
