"""Lightning module for training the MIS model."""

import os

import numpy as np
import scipy.sparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data

from co_datasets.mis_dataset import MISDataset
from pl_meta_model_rddm import COMetaModel
from utils.mis_utils import mis_decode_np


class MISModel(COMetaModel):
  def __init__(self,
               param_args=None):
    super(MISModel, self).__init__(param_args=param_args, node_feature_only=True)

    data_label_dir = None
    if self.args.training_split_label_dir is not None:
      data_label_dir = os.path.join(self.args.storage_path, self.args.training_split_label_dir)

    self.train_dataset = MISDataset(
        data_file=os.path.join(self.args.storage_path, self.args.training_split),
        data_label_dir=data_label_dir,
    )

    self.test_dataset = MISDataset(
        data_file=os.path.join(self.args.storage_path, self.args.test_split),
    )

    self.validation_dataset = MISDataset(
        data_file=os.path.join(self.args.storage_path, self.args.validation_split),
    )

  def forward(self, x, t, edge_index):
    return self.model(x, t, edge_index=edge_index)

  def generate_x_in(self, graph_data, point_indicator, edge_indicator):
    node_labels = graph_data.x

    x_in = torch.ones_like(node_labels, dtype=torch.float)
    x_in[::2] = 0

    return x_in

  def training_step(self, batch, batch_idx):
    exit(0)

  def categorical_denoise_step(self, xt, x_in_onehot, t, device, edge_index=None, target_t=None):
    with torch.no_grad():
      t = t.view(xt.shape[0])
      x_res_pred = self.forward(
          xt.float().to(device),
          t.float().to(device),
          edge_index.long().to(device) if edge_index is not None else None,
      )
      x_res_pred = x_res_pred.clamp(-1, 1)
      x0_pred_onehot = x_in_onehot - x_res_pred
      x0_pred_prob = x0_pred_onehot.reshape((1, xt.shape[0], -1, 2)).softmax(dim=-1)
      x_in_pred_prob = x_in_onehot.reshape((1, xt.shape[0], -1, 2)).softmax(dim=-1)

      xt = self.categorical_posterior(target_t.cpu(), t.cpu(), x0_pred_prob, xt, x_in_pred_prob)
      return xt

  def test_step(self, batch, batch_idx, draw=False, split='test'):
    device = batch[-1].device
    _, graph_data, point_indicator, edge_indicator = batch

    node_labels = graph_data.x
    edge_index = graph_data.edge_index

    stacked_predict_labels = []
    edge_index = edge_index.to(node_labels.device).reshape(2, -1)
    edge_index_np = edge_index.cpu().numpy()
    adj_mat = scipy.sparse.coo_matrix(
        (np.ones_like(edge_index_np[0]), (edge_index_np[0], edge_index_np[1])),
    )

    x_in_ori = self.generate_x_in(graph_data, point_indicator, edge_indicator)
    x_in_ori = x_in_ori.reshape(-1)
    for _ in range(self.args.sequential_sampling):
      x_in = x_in_ori.clone()
      x_in = x_in.float()
      noise = torch.randn_like(node_labels.float())
      if self.args.parallel_sampling > 1:
        noise = noise.repeat(self.args.parallel_sampling, 1, 1)
        x_in = x_in.repeat(self.args.parallel_sampling, 1)
        noise = noise.reshape(-1)
        noise = torch.randn_like(noise)

      x_in_tmp = x_in.clone()
      x_in_tmp = x_in_tmp.reshape(-1)
      x_in_onehot = F.one_hot(x_in_tmp.long(), num_classes=2).float()
      x_in = x_in * 2 - 1
      x_in = x_in * (1.0 + 0.05 * torch.rand_like(x_in))
      x_in = x_in.reshape(-1)
      xt = x_in + noise

      xt = (xt > 0).long()
      xt = xt.reshape(-1)

      if self.args.parallel_sampling > 1:
        edge_index = self.duplicate_edge_index(edge_index, node_labels.shape[0], device)

      steps = self.args.inference_diffusion_steps
      cur_time = torch.ones((xt.shape[0],), device=xt.device)
      step = 1.0 / steps
      for i in range(steps):
        s = torch.full((xt.shape[0], ), step, device=xt.device)
        if i == steps-1:
          s = cur_time
        t1 = cur_time
        t2 = cur_time - s
        xt = self.categorical_denoise_step(
            xt, x_in_onehot, t1, device, edge_index, target_t=t2)
        cur_time = cur_time - s

      predict_labels = xt.float().cpu().detach().numpy() + 1e-6
      stacked_predict_labels.append(predict_labels)

    predict_labels = np.concatenate(stacked_predict_labels, axis=0)
    all_sampling = self.args.sequential_sampling * self.args.parallel_sampling

    splitted_predict_labels = np.split(predict_labels, all_sampling)
    solved_solutions = [mis_decode_np(predict_labels, adj_mat) for predict_labels in splitted_predict_labels]
    solved_costs = [solved_solution.sum() for solved_solution in solved_solutions]
    best_solved_cost = np.max(solved_costs)

    gt_cost = node_labels.cpu().numpy().sum()
    gap = (gt_cost - best_solved_cost) / gt_cost * 100
    metrics = {
        f"{split}/gt_cost": gt_cost,
        f"{split}/gap(%)": gap,
    }
    for k, v in metrics.items():
      self.log(k, v, on_epoch=True, sync_dist=True)
    self.log(f"{split}/solved_cost", best_solved_cost, prog_bar=True, on_epoch=True, sync_dist=True)
    return metrics

  def validation_step(self, batch, batch_idx):
    return self.test_step(batch, batch_idx, split='val')
