from __future__ import annotations # In order to reference a class in annotations inside it

from typing import Union
from functools import wraps

import torch


class ParameterMismatch(Exception):
    '''Exception thrown when the `ModuleParameters` involved in an operation
    don't have the same names (keys in their disctionaries)'''
    def __init__(self, key):
        self._key = key
    
    def __str__(self):
        return 'One of the ModuleParameters does not have a parameter named {self._key}'


# Decorator that catches KeyError and raises ParameterMismatch
def _catch_key_error(f):
    @wraps(f)
    def wrapper(*args, **kwds):
        try:
            return f(*args, **kwds)
        except KeyError as e:
            raise ParameterMismatch(str(e)) from None
    
    return wrapper


class ModuleParameters:
    '''A class that allows treating the parameters of a module as a vector
    
    It stores the dict generated by the `named_parameters` generator of torch.nn.Module
    In order for operations between ModuleParameters to work their dictionaries need
    to have the same keys, and the coresponding parameters the same sizes.
    '''
    def __init__(self, module: torch.nn.Module, grad=False):
        if grad:
            self._parameters = {k: v.grad for k, v in module.named_parameters()}
        else:
            self._parameters = dict(module.named_parameters())
    
    @staticmethod
    def from_dict(dict):
        '''Initializes class from a dictionary'''
        params = ModuleParameters(torch.nn.Module()) # Create with empty Module
        params._parameters = {k: v for k, v in dict.items()}
        return params

    @_catch_key_error
    def __add__(self, other: ModuleParameters):
        return ModuleParameters.from_dict(
            {k: v + other._parameters[k] for k, v in self._parameters.items()}
        )
    
    @_catch_key_error
    def __iadd__(self, other: ModuleParameters):
        for k in self._parameters.keys():
            self._parameters[k] += other._parameters[k]
        return self

    @_catch_key_error
    def __sub__(self, other: ModuleParameters):
        return ModuleParameters.from_dict(
            {k: v - other._parameters[k] for k, v in self._parameters.items()}
        )

    @_catch_key_error
    def __isub__(self, other: ModuleParameters):
        for k in self._parameters.keys():
            self._parameters[k] -= other._parameters[k]
        return self

    @_catch_key_error
    def __mul__(self, other: Union[float, ModuleParameters]):
        if isinstance(other, ModuleParameters):
            return ModuleParameters.from_dict(
                {k: v * other._parameters[k] for k, v in self._parameters.items()}
            )
        else:
            return other * self

    @_catch_key_error
    def __imul__(self, other: Union[float, ModuleParameters]):
        if isinstance(other, ModuleParameters):
            for k in self._parameters.keys():
                self._parameters[k] *= other._parameters[k]
        else:
            for k in self._parameters.keys():
                self._parameters[k] *= other
        
        return self

    def __rmul__(self, other: float):
        return ModuleParameters.from_dict(
            {k: v * other for k, v in self._parameters.items()}
        )

    @_catch_key_error
    def __truediv__(self, other: Union[float, ModuleParameters]):
        if isinstance(other, ModuleParameters):
            return ModuleParameters.from_dict(
                {k: v / other._parameters[k] for k, v in self._parameters.items()}
            )
        else:
            return (1 / other) * self

    @_catch_key_error
    def __itruediv__(self, other: Union[float, ModuleParameters]):
        if isinstance(other, ModuleParameters):
            for k in self._parameters.keys():
                self._parameters[k] /= other._parameters[k]
        else:
            for k in self._parameters.keys():
                self._parameters[k] /= other
        
        return self

    def __pow__(self, other: float):
        return ModuleParameters.from_dict(
            {k: v ** other for k, v in self._parameters.items()}
        )
    
    def sum(self):
        s = 0.
        for v in self._parameters.values():
            s += v.sum()
        
        return s
    
    @_catch_key_error
    def dot(self, other: ModuleParameters):
        x = 0.
        for k,v in self._parameters.items():
            x += torch.dot(v.flatten(), other._parameters[k].flatten())
        
        return x

    def size(self):
        return sum(v.size() for v in self._parameters.values())

    def __repr__(self):
        return f'ModuleParameters:{repr(self._parameters)}'
