import json
import numpy as np
import yaml
import os

_ROOT = os.path.abspath(os.path.dirname(__file__))
def get_data(path):
    return os.path.join(_ROOT, 'data', path)
    
class RandomNames:
    def __init__(self):        
        self.prob_matrices = {}

    def generate_from_config_file(self, file=None):
        """
            This method calculate the probabities needed to 
            generate random names.
            Args:
                file: the yaml config file. In note the method use the default config file and name lists
        """
        returns = {}
        with open(file if file is not None else get_data("config.yml")) as f:
            config = yaml.load(f)
            for key,value in config.items():
                if file is None:
                    value = [get_data(x) for x in value]
                unique_letters, initial_prob_matrix, transition_prob_matrix = RandomNames._generate_prob(value)
                returns[key] = {"unique_letters":unique_letters,\
                                "initial_prob_matrix":initial_prob_matrix,\
                                "transition_prob_matrix":transition_prob_matrix}
        self.prob_matrices = returns
    
    def load_from_json(self,file):
        """
            Load the data needed from a json file
            file: a json file generated by the write_to_json method of this class
        """
        with open(file) as f:
            self.prob_matrices = json.load(f)

    def write_to_json(self,file):
        """
            Write the data used to generate random name to a json file. 
            The data can be restored using the load_from_json method of this class.

            Args:
                file: the path of the json file
        """
        with open(file,"w") as f:
            json.dump(self.prob_matrices,f)

    def get_type_names(self):
        """
            Return the types of name that can be generate by this class.
            This types depends on the config file
            
            Return:
                An array containing the type of names.
        """
        return list(self.prob_matrices.keys())

    def generate_random_name(self,name_type, length, init = ""):
        """
            This method is used to generate random names.
            
            Args:
                name_type: the type of name generated
                length: the legth of the name
                init: This parameter must be used if the name must start with a specific combination of characters.
            Return:
                A random name 
        """
        if name_type not in self.prob_matrices.keys():
            raise KeyError("name_type %s not in prob_matrices"%(name_type,))
        if type(length) != int:
            raise TypeError("lenght must be integer")
        if length <= 0:
            raise ValueError("length must be greater than 0")
        
        if init != "":
            ret_string = init.lower()
        else:
            ret_string = np.random.choice(self.prob_matrices[name_type]["unique_letters"],\
                                p=self.prob_matrices[name_type]["initial_prob_matrix"])
        
        for _ in range(len(ret_string),length):
            ret_string += np.random.choice(self.prob_matrices[name_type]["unique_letters"],\
                                p=self.prob_matrices[name_type]["transition_prob_matrix"]\
                                    [self.prob_matrices[name_type]["unique_letters"].index(ret_string[-1])])
        return ret_string.capitalize()
        
    
    

    @staticmethod
    def _generate_prob(files):
        """
            Given a list of file in json format containing an array of names
            this static method learn 3 things:
            - The list of unique charachters
            - An array containing the probabilities for each character to be the first of a name
            - A matrix containig the transition probabilities from one character to another
        """
        transition_prob = {}
        initial_prob = {}
        unique_letters = set()

        for file in files:
            with open(file,encoding="utf-8") as f:
                data = json.load(f)
                for name in data:
                    name = name.lower()
                    try:
                        initial_prob[name[0]] += 1
                    except:
                        initial_prob[name[0]] = 1
                    for idx in range(1,len(name)):
                        if name[idx-1] not in transition_prob.keys():
                            transition_prob[name[idx-1]] = {}
                        if name[idx] not in transition_prob[name[idx-1]].keys():
                            transition_prob[name[idx-1]][name[idx]] = 0
                        transition_prob[name[idx-1]][name[idx]] += 1
        
        unique_letters = unique_letters.union(set(initial_prob.keys()))
        unique_letters = unique_letters.union(set(transition_prob.keys()))
        for key in transition_prob.keys():
            unique_letters = unique_letters.union(set(transition_prob[key].keys()))
        unique_letters = list(unique_letters)
        
        initial_prob_matrix = np.zeros(len(unique_letters))
        transition_prob_matrix = np.zeros((len(unique_letters),len(unique_letters)))
        for idx, char in enumerate(unique_letters):
            try:
                initial_prob_matrix[idx] = initial_prob[char]
            except:
                pass
            for idx2, char2 in enumerate(unique_letters):
                try:
                    transition_prob_matrix[idx][idx2] = transition_prob[char][char2]
                except:
                    pass
        initial_prob_matrix /= initial_prob_matrix.sum()
        transition_prob_matrix = (transition_prob_matrix.T/transition_prob_matrix.sum(axis=1)).T
        return unique_letters, initial_prob_matrix.tolist(), transition_prob_matrix.tolist()

