#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""

YAVUZ ÇETİN ÖLÜMSÜZDÜR

@author: Mehmet
"""

import networkx as nx
import pandas as pd
from networkx.drawing.nx_agraph import graphviz_layout, to_agraph
"""
Model class for the Markov Chains object
"""
class Chain:

    """
    Takes the sample name and a list containing the names of the instructions

    Creates a graph G, which is a Directed MultiGraph networkx object and M,
    which is initially zero but after calling adjacency_matrix_to_pd_dataframe()
    it contains the adjacency matrix

    Notice that instruction_names should contain unique names
    """
    def __init__(self, name, instruction_names):
        self.name = name
        self.instruction_names = instruction_names
        self.num_of_instructions = len(instruction_names)
        self.G = nx.MultiDiGraph()
        self.M = None

        for i in range(self.num_of_instructions):
            self.G.add_node(instruction_names[i])


    """
    Adds a single instance to the graph
    """
    def add(self, source_instr, target_instr, weight):
        self.G.add_edges_from([(source_instr, target_instr, {'weight': weight})])


    """
    Takes a list object that contains tuples in the format (source_instr, source_target, weight)
    """
    def add_multiple(self, _list):
        to_add = [(w[0], w[1], {'weight': w[2]}) for w in _list]
        self.G.add_edges_from(to_add)


    """
    Returns the adjacency matrix
    """
    def adjacency_matrix_to_pd_dataframe(self):
        self.M = nx.to_pandas_adjacency(self.G)
        return self.M

    """
    Check if the adjacency matrix satisfied the condition where
    probabilities of all events originating from a single node sums to 1
    """
    def check_validity(self):
        if self.M is None:
            raise Exception('Adjacency matrix is not calculated yet!')
        else:
            sum_vector = self.M.sum(axis=1)
            unique = list(set(sum_vector))
            is_valid = True

            for i in range(len(unique)):
                if unique[i] not in (0, 1):
                    is_valid = False

        return is_valid

    """
    !!!!!!!! BIG TODO !!!!!!!!!
    Change the pickle serialization into h5py
    """
    """
    Save the model (which is, in fact, the adjacency matrix)
    """
    def save_model(self, path='chain_model.pkl'):
        self.M.to_pickle(path)


    """
    Load the model from a pickle file
    """
    def load_model(self, path):
        self.M = pd.read_pickle(path)


    """
    Draws the graph into a specified path
    """
    def draw(self, draw_path='graph.png'):
        # draw the graph with graphviz
        # get weights
        for u,v,d in self.G.edges(data=True):
            d['label'] = d.get('weight','')
        # set defaults
        A = to_agraph(self.G)
        A.layout('dot')
        A.draw('graph.png')


###### testing
test_instruction_data = ['pop', 'push', 'move', 'add', 'sub', 'add', 'or']

chain = Chain('zeus', test_instruction_data)
chain.add_multiple([('pop', 'move', 0.4),
                  ('pop', 'add', 0.4),
                  ('pop', 'pop', 0.2)])

chain.add('or', 'pop', 0.4)
chain.add('or', 'move', 0.6)


M = chain.adjacency_matrix_to_pd_dataframe()
chain.draw()