import argparse
import itertools
from functools import partial

import numpy as np
import onnx
import torch
from tqdm import tqdm

from data_processing import datasets
from pointnet.segmentation_model import PointNetSegmentation
from relaxations import taylor
from relaxations.interval import Interval
from relaxations.refine_intervals import split_implicitly
from transformations.transformations import parse_transformation
from util import onnx_converter
from util.argparse import parse_theta, absolute_path
from util.experiment import Experiment
from util.math import set_random_seed, logits_to_category, DEFAULT_SEED
from util.timing import Timer
from verifier.eran_verifier import EranVerifier

parser = argparse.ArgumentParser()
parser.add_argument('--model', type=absolute_path, required=True, help="Path to the model to verify (.pth)")
parser.add_argument('--num_points', type=int, default=1024, help="Number of points per point cloud")
parser.add_argument('--experiment', type=str, help="Unique identifier name for the experiment. Will be generated by default")
parser.add_argument('--theta', type=parse_theta, nargs='+', required=True, help="List of transformation parameters to certify. Either number or number followed by 'deg' for degree.")
parser.add_argument('--intervals', type=int, default=1, help="Number of parameter splits per dimension")
parser.add_argument('--transformation', type=str, default='RotationZ', help="The transformation function, or composition of multiple transformations separated by '+'")
parser.add_argument('--pooling', choices=['improved_max'], default='improved_max', help='The pooling function to use')
parser.add_argument('--relaxation', default='taylor', choices=['taylor', 'interval'], help="The relaxation to use")
parser.add_argument('--seed', type=int, default=DEFAULT_SEED, help='Seed for random number generators')
parser.add_argument('--implicit-intervals', type=int, default=1, help="Number of implicit splits for interval bound refinement")

settings = parser.parse_args()

experiment = Experiment(settings)
logger = experiment.logger
checkpoints = experiment.load_checkpoints()

transformation = parse_transformation(settings.transformation)
parameters = settings.theta
if len(parameters) != transformation.num_params:
    raise ValueError(f"The number of parameters does no match those required by the transformation. Expected {transformation.num_params}, received {parameters}")

test_data = datasets.shapenet(num_points=settings.num_points, split='test', rotate='none')
num_total_classes = 50

torch_model = PointNetSegmentation(
    number_points=settings.num_points,
    num_seg_classes=num_total_classes,
    encode_onnx=True
)
torch_model.load_state_dict(torch.load(settings.model, map_location=torch.device('cpu')))
torch_model = torch_model.eval()

export_file = settings.model.with_suffix('.onnx')
onnx_model = onnx_converter.convert(torch_model, settings.num_points, export_file)
logger.info(onnx.helper.printable_graph(onnx_model.graph))

eran = EranVerifier(onnx_model)
test_samples = len(test_data)
interval = len(test_data) // 100

total_percentage_certified = 0.0
iterations = 0

timer = Timer()

for counter, i in enumerate(range(0, test_samples, interval)):

    iterations += 1
    np_points, label = test_data[i]
    points = torch.from_numpy(np_points)
    points = torch.unsqueeze(points, dim=0)

    prediction = torch_model(points.transpose(2, 1)).detach().cpu().numpy()[0]
    predicted_parts = logits_to_category(prediction, label)

    correct_predicted_indexes = np.argwhere(predicted_parts == label).squeeze()
    correct_predicted_labels = label[correct_predicted_indexes]

    num_params = transformation.num_params
    intervals_per_dimension = [settings.intervals] * num_params

    theta_inc = [2.0 * param / float(intervals) for param, intervals in zip(parameters, intervals_per_dimension)]

    checkpoints_sample = checkpoints[str(i)] if str(i) in checkpoints else {}

    valid_classes = np.unique(label)
    certified_points = np.full(settings.num_points, True)

    timer.start()

    parameter_iterator = itertools.product(*(range(intervals) for intervals in intervals_per_dimension))
    progress_bar = tqdm(parameter_iterator, desc=f"Object {counter}", unit="interval", total=settings.intervals ** num_params)
    for idx in progress_bar:
        assert len(idx) == len(parameters) and len(theta_inc) == len(parameters)
        params = [Interval(j * inc - theta, (j + 1) * inc - theta) for j, inc, theta in zip(idx, theta_inc, parameters)]

        interval_key = 'x'.join([f"[{i.lower_bound:.4f},{i.upper_bound:.4f}]" for i in params])
        if interval_key in checkpoints_sample:
            interval_certified = np.array(checkpoints_sample[interval_key])

        else:
            assert np_points.shape[0] == settings.num_points, \
                f"invalid points shape {np_points.shape}, expected ({settings.num_points}, x)"

            if settings.relaxation == 'interval':
                bounds = transformation.transform(np_points, params)

                if settings.implicit_intervals > 1:
                    bounds = split_implicitly(
                        bounds=bounds, intervals=settings.implicit_intervals,
                        encode_rotation=partial(transformation.transform, np_points),
                        params=params
                    )

                (dominant_classes, nlb, nub) = eran.analyze_segmentation_box(bounds, label, valid_classes, num_total_classes)

            elif settings.relaxation == 'taylor':
                bounds = transformation.transform(np_points, params)
                constraints = taylor.encode(transformation, np_points, params)

                if settings.implicit_intervals > 1:
                    bounds = split_implicitly(
                        bounds=bounds, intervals=settings.implicit_intervals,
                        encode_rotation=partial(transformation.transform, np_points),
                        params=params
                    )

                (dominant_classes, nlb, nub) = eran.analyze_segmentation_linear(bounds, constraints, params, label, valid_classes, num_total_classes)
                dominant_classes = np.array(dominant_classes)

                assert np.all(np.logical_or(label == dominant_classes, dominant_classes == -1)), \
                    f"Wrong dominant class! label {label}, dominant_class: {dominant_classes}"

            else:
                raise Exception(f"Invalid relaxation {settings.relaxation}")

            interval_certified = dominant_classes == label
            checkpoints_sample[interval_key] = interval_certified.tolist()

        certified_points = np.logical_and(certified_points, interval_certified)

    elapsed = timer.stop()

    checkpoints[str(i)] = checkpoints_sample
    experiment.store_checkpoints(checkpoints)

    correct_predicted_certified = certified_points[correct_predicted_indexes]

    percentage_certified = np.mean(correct_predicted_certified.astype(float))
    total_percentage_certified += percentage_certified

    logger.info(f"Successfully certified {percentage_certified} of points in {settings.intervals} intervals")

    logger.info(f"Time for this round: {elapsed}s. Total time: {timer.get()}s.")
    logger.info(f"Tested {iterations} data points for which on average {total_percentage_certified / iterations} of points were correctly predicted.")
