import torch
from torch.utils.data import DataLoader, TensorDataset
from dataclasses import dataclass

@dataclass
class GD_output:
    parameters: dict = None
    objectives: torch.Tensor = None
    per_step_objectives: torch.Tensor = None
    per_epoch_objectives: torch.Tensor = None
    epoch_step_nums: torch.Tensor = None
    grad_steps: range = None
    lr: float = None
    num_steps: int = None
    decay_rate: float = None
    batch_size: int = None
    num_epochs: int = None
    max_steps: int = None
    type_flag: str = None
    
    def __str__(self):
        print_string = ''
        if self.type_flag == 'gd':
            print_string += 'GRADIENT DESCENT OUTPUT'
            print_string += f'\nlearning rate: {self.lr}, decay rate: {self.decay_rate}, gradient steps: {self.num_steps}'
            for key, parameter in self.parameters.items():
                print_string += f'\n\nparameter {key}:\n{parameter}'
            print_string += f'\n\nobjectives:\n{self.objectives}'
            return print_string
        else:
            print_string += 'STOCHASTIC GRADIENT DESCENT OUTPUT'
            print_string += f'\nlearning rate: {self.lr}, decay rate: {self.decay_rate},' \
                f' batch size: {self.batch_size}, number of epochs: {self.num_epochs}'
            for key, parameter in self.parameters.items():
                print_string += f'\n\nparameter {key}:\n{parameter}'
            print_string += f'\n\nper-step objectives:\n{self.per_step_objectives}'
            print_string += f'\n\nper-epoch mean objectives:\n{self.per_epoch_objectives}'
            print_string += f'\n\nepochs began/completed on the follow gradient steps:\n{self.epoch_step_nums}'
            return print_string


def GD(J, init_parameters, lr, num_steps, decay_rate=0):

    single_parameter = False
    if isinstance(init_parameters, torch.Tensor):
        init_parameters = {'theta': init_parameters}
        single_parameter = True
        
    parameters = {key: tensor.clone().requires_grad_() for key, tensor in init_parameters.items()}
    input_parameters = list(parameters.values())[0] if single_parameter else parameters
    objective = J(input_parameters)
    objectives = [objective.detach()]
    running_parameters = {key: [tensor.clone()] for key, tensor in init_parameters.items()}

    for t in range(num_steps):

        objective.backward()

        with torch.no_grad():
            for parameter in parameters.values():
                parameter -= lr * (1 - decay_rate) ** (1 + t) * parameter.grad

        input_parameters = list(parameters.values())[0] if single_parameter else parameters
        objective = J(input_parameters)

        objectives.append(objective.detach())
        for key in running_parameters.keys():
                running_parameters[key].append(parameters[key].detach().clone())
        for parameter in parameters.values():
            parameter.grad.zero_()

    for key in running_parameters.keys():
        running_parameters[key] = torch.row_stack(running_parameters[key])
    objectives = torch.stack(objectives)
    output = GD_output(parameters=running_parameters,
                       objectives=objectives,
                       lr=lr,
                       num_steps=num_steps,
                       decay_rate=decay_rate,
                       type_flag='gd')
    return output


def SGD(g, init_parameters, X, lr, batch_size, num_epochs, y=None, decay_rate=0, max_steps=-1, shuffle=True, random_state=None):

    if random_state != None:
        torch.manual_seed(random_state)
    dataset = TensorDataset(X, y) if y != None else X
    data_loader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=shuffle)
    num_mini_batches = len(data_loader)
    
    single_parameter = False
    if isinstance(init_parameters, torch.Tensor):
        init_parameters = {'theta': init_parameters}
        single_parameter = True
    
    parameters = {key: tensor.clone().requires_grad_() for key, tensor in init_parameters.items()}
    per_step_objectives = []
    per_epoch_objectives = []
    running_parameters = {key: [tensor.clone()] for key, tensor in init_parameters.items()}
    first_step = True
    s = 0
    epoch_step_nums = [0]

    for t in range(num_epochs):
        for i, mini_batch in enumerate(data_loader):
            
            input_parameters = list(parameters.values())[0] if single_parameter else parameters

            if first_step:
                objective = g(*mini_batch, input_parameters).mean() if y != None else g(mini_batch, input_parameters).mean()
                per_step_objectives.append(objective.detach())
                per_epoch_objectives.append(objective.detach())
                first_step = False

            objective.backward()

            with torch.no_grad():
                for parameter in parameters.values():
                    parameter -= lr * (1 - decay_rate) ** (1 + s) * parameter.grad
            
            objective = g(*mini_batch, input_parameters).mean() if y != None else g(mini_batch, input_parameters).mean()

            per_step_objectives.append(objective.detach())
            for key in running_parameters.keys():
                running_parameters[key].append(parameters[key].detach().clone())
            for parameter in parameters.values():
                parameter.grad.zero_()

            complete_epoch = True if i + 1 == num_mini_batches else False
            s += 1
            if s == max_steps:
                break
        
        if complete_epoch:
            epoch_step_nums.append(s)
        
        epoch_objective = torch.tensor(per_step_objectives[num_mini_batches * t + 1:])
        per_epoch_objectives.append(epoch_objective.mean())

        if s == max_steps:
            break

    for key in running_parameters.keys():
        running_parameters[key] = torch.stack(running_parameters[key], dim=0)
    per_step_objectives = torch.stack(per_step_objectives)
    per_epoch_objectives = torch.stack(per_epoch_objectives)
    epoch_step_nums = torch.tensor(epoch_step_nums)
    output = GD_output(parameters=running_parameters,
                       per_step_objectives=per_step_objectives,
                       per_epoch_objectives=per_epoch_objectives,
                       epoch_step_nums=epoch_step_nums,
                       grad_steps=range(len(per_step_objectives)),
                       lr=lr,
                       num_epochs=num_epochs,
                       decay_rate=decay_rate,
                       batch_size=batch_size,
                       max_steps=max_steps)
    return output