###########################################################
#                                                         #
#    This file contains various classes and methods       #
#    needed to handle training systems using either       #
#    BDD or ROBDD methods                                 #
#                                                         #
###########################################################

from c_bdd import BDDStructure
from c_robdd import ROBDDStructure

import pandas as pd
import math


class BDDTrain(BDDStructure):
    def __init__(self, df):
        super().__init__(df.shape[1] - 1)
        self.cols = list(df.columns)
        self.reliability = df[self.cols[0]].mean()
        self.dataFrame = df
        self.dominantColumn = self.findDominantColumn()
        
    def findDominantColumn(self):
        max_i = -1
        max_c = 0.0
        for i in range(self.order):
                covariance = self.dataFrame[self.cols[0]].cov(self.dataFrame[self.cols[i+1]])
                if covariance > max_c:
                    max_i = i
                    max_c = covariance
        return max_i
        
    def selectPivotIndex(self):
        self.pivotIndex = max(0, self.dominantColumn)
    
    def restriction(self):
        df = self.dataFrame[self.dataFrame[self.cols[self.pivotIndex+1]] == 0]
        df.pop(self.cols[self.pivotIndex+1])
        struct = BDDTrain(df)
        for i in range(self.order):
            if i != self.pivotIndex:
                struct.addComponent(self.components[i])
        return struct
    
    def contraction(self):
        df = self.dataFrame[self.dataFrame[self.cols[self.pivotIndex+1]] == 1]
        df.pop(self.cols[self.pivotIndex+1])
        struct = BDDTrain(df)
        for i in range(self.order):
            if i != self.pivotIndex:
                struct.addComponent(self.components[i])
        return struct
   
    def isFunctioning(self):
        if self.dominantColumn >= 0:
            return (self.reliability == 1)
        else:
            return (self.reliability > 0.5)
    
    def isFailed(self):
        if self.dominantColumn >= 0:
            return (self.reliability == 0)
        else:
            return (self.reliability <= 0.5)
           
    


# Reliability tolerance level
EPSILON = 0.01

def SET_EPSILON(eps):
    global EPSILON
    EPSILON = eps 

# If MERGE_STRUCTURES is True, dataframes will be merged in the simplification process
MERGE_STRUCTURES = True

def SET_MERGE_STRUCTURES(flag):
    global MERGE_STRUCTURES
    MERGE_STRUCTURES = flag


class ROBDDTrain(ROBDDStructure):
    def __init__(self, df):
        super().__init__(df.shape[1] - 1)
        self.cols = list(df.columns)
        self.reliability = df[self.cols[0]].mean()
        self.dataFrame = df
        
    def restriction(self):
        df = self.dataFrame[self.dataFrame[self.cols[1]] == 0]
        df.pop(self.cols[1])
        return ROBDDTrain(df)
    
    def contraction(self):
        df = self.dataFrame[self.dataFrame[self.cols[1]] == 1]
        df.pop(self.cols[1])
        return ROBDDTrain(df)
   
    def isFunctioning(self):
        return (self.reliability == 1)
    
    def isFailed(self):
        return (self.reliability == 0)
    
    def isEquivalentTo(self, s):
        diff = math.fabs(self.reliability - s.reliability)
        return (diff < EPSILON)
    
    def mergeStructure(self, s):
        if MERGE_STRUCTURES:
            dfs = [self.dataFrame, s.dataFrame]
            self.dataFrame = pd.concat(dfs)