import argparse
import itertools
from functools import partial

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

from data_processing import datasets
from pointnet.model import PointNet
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 DEFAULT_SEED
from util.timing import Timer, BOUND_TIMER, TOTAL_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('--implicit-intervals', type=int, default=1, help="Number of implicit splits for interval bound refinement")
parser.add_argument('--transformation', type=str, default='RotationZ', help="The transformation function, or composition of multiple transformations separated by '+'")
parser.add_argument('--relaxation', default='taylor', choices=['taylor', 'interval'], help="The relaxation to use")
parser.add_argument('--pooling', choices=['improved_max', 'max', 'avg'], default='improved_max', help='The pooling function to use')
parser.add_argument('--max_features', type=int, default=1024, help='The number of global features')
parser.add_argument('--seed', type=int, default=DEFAULT_SEED, help='seed for random number generator')

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.modelnet40(num_points=settings.num_points, split='test', rotate='none')

torch_model = PointNet(
    number_points=settings.num_points,
    num_classes=test_data.num_classes,
    max_features=settings.max_features,
    pool_function=settings.pooling,
    disable_assertions=True,
    transposed_input=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

correct_predictions = 0
verified_same = 0
verified_different = 0
not_verified = 0
iterations = 0

timer = Timer()

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

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

    prediction = torch_model(points.transpose(2, 1))
    max_prediction = prediction.data.max(1)[1].item()
    correct = max_prediction == label

    if not correct:
        logger.info(f"Incorrect prediction, skipping. True label was {label}, prediction was {max_prediction}")
        continue

    correct_predictions += 1

    session = onnxruntime.InferenceSession(onnx_model.SerializeToString())
    input_data = np.expand_dims(np.transpose(np_points.copy()), axis=(0, -1))
    inputs = {session.get_inputs()[0].name: input_data}
    outputs = session.run(None, inputs)

    np.testing.assert_allclose(prediction.detach().numpy(), outputs[0], rtol=1e-2, atol=1e-3)

    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 {}

    certified = True

    timer.start(TOTAL_TIMER)
    elapsed_bounds = 0.0

    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 = 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':
                timer.start(BOUND_TIMER)
                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
                    )
                elapsed_bounds += timer.stop(BOUND_TIMER)

                (dominant_class, nlb, nub) = eran.analyze_classification_box(bounds)

            elif settings.relaxation == 'taylor':
                timer.start(BOUND_TIMER)
                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
                    )
                elapsed_bounds += timer.stop(BOUND_TIMER)

                (dominant_class, nlb, nub) = eran.analyze_classification_linear(bounds, constraints, params)

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

            interval_certified = dominant_class == label.item()
            checkpoints_sample[interval_key] = interval_certified

        if not interval_certified:
            certified = False

    elapsed_total = timer.stop(TOTAL_TIMER)

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

    if certified:
        logger.info(f"Successfully certified class {label.item()} in {settings.intervals} intervals")
        verified_same += 1
    else:
        logger.info(f"Failed to certify class {label.item()}")

    logger.info(f"Time for this round: (bounds: {elapsed_bounds}s) (total: {elapsed_total}s). Total time: (bounds: {timer.get(BOUND_TIMER, 0.0)}s) (total: {timer.get(TOTAL_TIMER)}s).")
    logger.info(f"Tested {iterations} data points out of which {correct_predictions} were correctly predicted.")
    logger.info(f"Successfully certified {verified_same} samples.")
