import numpy as np
from radnn import mlsys, FileStore

# -----------------------------
# Standard Libraries
# -----------------------------
import time

# -----------------------------
# PyTorch
# -----------------------------
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler

from radnn.evaluation import EvaluateClassification
from radnn.plots import PlotConfusionMatrix, PlotLearningCurve
from radnn.experiment import MLExperimentLog, experiment_fold_number, experiment_name_with_fold
import matplotlib.pyplot as plt
from .lr_schedulers import StairCaseLR, ErrorPlateauWithWarmup, AccuracyPlateauWithWarmup, CompositeScheduler
from radnn.errors import *

# -----------------------------
# Progress Bar
# -----------------------------
from tqdm import tqdm





class GradientDescentMethod(object):
  # --------------------------------------------------------------------------------------------------------------------
  def __init__(self, model, hyperparams):
    self.hprm = hyperparams
    self._descr = ""
    
    self.is_based_on_metrics = False
    
    self.optimizer = None
    self.lr_scheduler = None
    self.is_monitoring_minimization = True
    self.build_optimizer(model)
    self.build_lr_scheduler(model)
    
  # --------------------------------------------------------------------------------------------------------------------
  @property
  def last_lr(self):
    return self.lr_scheduler.get_last_lr()[0]
  # --------------------------------------------------------------------------------------------------------------------
  def __str__(self):
    return self._descr
  # --------------------------------------------------------------------------------------------------------------------
  def next_step(self, metrics=None):
    nLR = None
    if self.lr_scheduler is not None:
      oScheduler = self.lr_scheduler

      if isinstance(self.lr_scheduler, CompositeScheduler):
        oCompScheduler: CompositeScheduler = self.lr_scheduler
        oCompScheduler.step(metrics)
      else:
        if self.is_based_on_metrics:
          oScheduler.step(metrics)
        else:
          oScheduler.step()
      nLR = self.lr_scheduler.get_last_lr()[0]
      
    return nLR
  
  # --------------------------------------------------------------------------------------------------------------
  def exclude_params_from_weight_decay(self, model, weight_decay):
    oWeightDecayParams = []
    oUnlimitedParams = []
    for name, p in model.named_parameters():
      if not p.requires_grad:
        continue
      
      if "temperature" in name:
        oUnlimitedParams.append((name,p))
      else:
        oWeightDecayParams.append((name,p))
    if len(oUnlimitedParams) == 0:
      oGroups = model.named_parameters()
    else:
      oGroups = (
        {"params": oWeightDecayParams},
        {"params": oUnlimitedParams, "weight_decay": 0.0},
      
      )
    return oGroups
  # --------------------------------------------------------------------------------------------------------------------
  def build_optimizer(self, model):
    hprm = self.hprm
    nLR = hprm["Training.LearningRate"]
    nWeightDecay = hprm.get("Training.Regularize.WeightDecay", float(0))
    sExtra = ""
    oParamsToOptimize = self.exclude_params_from_weight_decay(model, nWeightDecay)
    
    if hprm["Training.Optimizer"].upper() == "SGD":
      self.optimizer = optim.SGD(oParamsToOptimize, lr=nLR,
                                 momentum=hprm.get("Training.Momentum", 0.0),
                                 nesterov=hprm.get("Training.Momentum.Nesterov", False),
                                 weight_decay=nWeightDecay)
      sExtra = f'momentum={self.optimizer.defaults["momentum"]}'
      if self.optimizer.defaults["nesterov"]:
        sExtra += " (Nesterov)"
    elif hprm["Training.Optimizer"].upper() == "RMSPROP":
      self.optimizer = optim.RMSprop(oParamsToOptimize, lr=nLR,
                                     weight_decay=nWeightDecay,
                                     momentum=hprm.get("Training.Momentum", 0.0),
                                     eps=hprm.get("Training.RMSProp.Epsilon", 1e-8)
                                     )
    elif hprm["Training.Optimizer"].upper() == "ADAM":
      self.optimizer = optim.Adam(oParamsToOptimize, lr=nLR,
                                  weight_decay=nWeightDecay)
    elif hprm["Training.Optimizer"].upper() == "ADAMW":
      self.optimizer = optim.AdamW(oParamsToOptimize, lr=nLR,
                                   weight_decay=nWeightDecay)
    
    self._descr = f'Using {hprm["Training.Optimizer"].upper()} optimizer {sExtra}'
  # --------------------------------------------------------------------------------------------------------------------
  def build_lr_scheduler(self, model):
    hprm = self.hprm
    nLearningRate = hprm["Training.LearningRate"]
    sSchedulingType = hprm.get("Training.LearningRateSchedule", "MultiStepDivisor")
    oMilestoneEpochs: list = hprm.get("Training.LearningRateSchedule.Epochs", [])
    nChangeFactor = hprm.get("Training.LearningRateSchedule.ChangeFactor",
                             hprm.get("Training.LearningRateSchedule.StepRatio", None))
    oWaypointSetup = oMilestoneEpochs
    if np.asarray(oMilestoneEpochs).ndim == 1:
      nLR = nLearningRate
      oWaypointSetup = [[0,  nLR]]
      for nEpoch in oMilestoneEpochs:
        nLR = nLR * nChangeFactor
        oWaypointSetup.append([nEpoch, nLR])
    
    if len(oMilestoneEpochs) > 0:
      if oMilestoneEpochs[0][0] != 0:
        oMilestoneEpochs.insert(0, [0,  nLearningRate])
    
    nFinalChangeEpoch, nFinalLR = oWaypointSetup[-1]
    
    sExtra = ""
    self.lr_scheduler = None
    if (sSchedulingType.upper() == "MultiStepDivisor".upper()):
      if "Training.LearningRateSchedule.Epochs" in hprm:
        assert nChangeFactor is not None, TRAINER_LR_SCHEDULER_CHANGE_FACTOR_NOT_SET
        sExtra = f'milestones={oMilestoneEpochs}'
        self.lr_scheduler = lr_scheduler.MultiStepLR(self.optimizer,
                                                     milestones=oMilestoneEpochs,
                                                     gamma=nChangeFactor)
        sExtra = f'schedule={hprm["Training.LearningRateSchedule.Epochs"]} change_factor={nChangeFactor}'
      else:
        raise Exception(TRAINER_LR_SCHEDULER_INVALID_MILESTONE_SETUP)
    elif (sSchedulingType.upper() == "Staircase".upper()):
      if np.asarray(oWaypointSetup).ndim == 2:
        sExtra = f'milestones={oWaypointSetup}'
        if not isinstance(oWaypointSetup, list):
          raise Exception(TRAINER_LR_SCHEDULER_INVALID_SETUP)
        self.lr_scheduler = StairCaseLR(self.optimizer, oWaypointSetup)
      else:
        raise Exception(TRAINER_LR_SCHEDULER_INVALID_SETUP)
    
    elif sSchedulingType.upper() == "CosineAnnealing".upper():
      sExtra = f"T_max={nFinalChangeEpoch} final_lr={nFinalLR}"
      self.lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(self.optimizer,
                                                                     T_max=nFinalChangeEpoch,
                                                                     eta_min=nFinalLR)
    elif sSchedulingType.upper() == "ErrorPlateauWarmup".upper():
      assert nChangeFactor is not None, TRAINER_LR_SCHEDULER_CHANGE_FACTOR_NOT_SET
      dSpecialHprm = hprm["Training.LearningRateSchedule.ErrorPlateauWarmup"]
      self.lr_scheduler = ErrorPlateauWithWarmup(self.optimizer,
                                                   warmup_epochs = dSpecialHprm.get("WarmupEpochs", 2),
                                                   warmup_lr = dSpecialHprm.get("WarmupLearningRate", 0.1),
                                                   base_lr = nLearningRate,
                                                   change_factor=nChangeFactor,
                                                   patience = dSpecialHprm.get("Patience", 3)
                                                 )
      sExtra = f'Warmup: epochs={dSpecialHprm.get("WarmupEpochs", 2)} lr={dSpecialHprm.get("WarmupLearningRate", 0.1)} change_factor={nChangeFactor} monitoring minimization metric'
    elif sSchedulingType.upper() == "AccuracyPlateauWarmup".upper():
      self.is_monitoring_minimization = False
      assert nChangeFactor is not None, TRAINER_LR_SCHEDULER_CHANGE_FACTOR_NOT_SET
      dSpecialHprm = hprm["Training.LearningRateSchedule.AccuracyPlateauWarmup"]
      self.lr_scheduler = AccuracyPlateauWithWarmup(self.optimizer,
                                                   warmup_epochs = dSpecialHprm.get("WarmupEpochs", 2),
                                                   warmup_lr = dSpecialHprm.get("WarmupLearningRate", 0.1),
                                                   base_lr = nLearningRate,
                                                   change_factor=nChangeFactor,
                                                   patience = dSpecialHprm.get("Patience", 3)
                                                 )
      sExtra = f'Warmup: epochs={dSpecialHprm.get("WarmupEpochs", 2)} lr={dSpecialHprm.get("WarmupLearningRate", 0.1)} change_factor={nChangeFactor} monitoring maximization metric'
    assert self.lr_scheduler is not None, TRAINER_LR_SCHEDULER_UNSUPPORTED % sSchedulingType
    self._descr += f" with scheduler {sSchedulingType} " + sExtra