import pandas as pd
import utilities as utils
import progressbar as pb
import logging

### --------------------------
### Define methods - data prep
### --------------------------

def prep_data_pairs(in_file):
    """Read time series data from file and get *obs_pair_hash* and *alphabet*.
    
    Parameters:
        in_file: time series data (csv format) of form var,time.

    Returns in the following order:
        obs_pair_hash:  dictionary | key: time t | val: list of discrete or discretized vars occurring at t
        obs_times_hash:  dictionary | key: variable x | val: list of time points where x ocurred.
        alphabet:       list | discrete or discretized vars in the time series
        events:         number of timepoints in the time series
        max_time:       maximum time stamp
    """

    logging.log(logging.DEBUG, "Prepping data structure.")

    obs_pair_hash = {}
    obs_times_hash = {}
    alphabet = []
    events = 0
    max_time = 0

    df = pd.read_csv(in_file, header=None)
    events = len(df.index)

    for index, row in df.iterrows():
        case, val, t = row
        
        # Populate obs_pair_hash
        if t not in obs_pair_hash:
            obs_pair_hash[t] = [(case, val)]
        else:
            obs_pair_hash[t].append((case, val))

        # Populate obs_times_hash : continue with building the subdictionaries where each variable is populated with a dictionary of time stamps containing a list of their cases
        if val not in obs_times_hash:
            obs_times_hash[val] = {t: [case]}
        else:
            if t in obs_times_hash[val]:
                obs_times_hash[val][t].append(case)
            else:
                obs_times_hash[val][t] = [case]

        if val not in alphabet:
            alphabet.append(str(val))
    
    max_time = df[2].values[-1]

    return(obs_pair_hash, obs_times_hash, alphabet, events, max_time)

### -----------------------------------
### Define methods - prima facie causes
### -----------------------------------

def generate_hypotheses_for_effects(cause_alpha, effects_alpha, window_start, window_end):
    """
    Generates hypotheses for all effects. A hypothesis is of form:
        (cause effect window-start window-end)

    Parameters:
        cause_alpha:    a variable or set of variables
        effects_alpha:  a list of possible effects
        window_start:   the start time of a window, i.e. r in window [r,s]
        window_end:     the end time of a window, i.e. s in window [r,s]

    Returns:
        A list of tuples containing the hypotheses for possible prima facie causes
    """

    logging.log(logging.DEBUG, "Generating hypotheses w.r.t. prima facie causes.")

    hypotheses = []
    for e in effects_alpha:
        hypotheses.extend(generate_hypotheses_for_effect(cause_alpha, e, window_start, window_end))
    
    return(hypotheses)

def generate_hypotheses_for_effect(cause_alpha, effect, window_start, window_end):
    """
    Lists all hypotheses for a given effect. Excludes hypotheses where the cause and effect are the same variable.
    
    See the docs of `generate_hypotheses_for_effects`
    """
    hyps = []
    for cause in cause_alpha:
        if cause != effect:
            t = (cause, effect, window_start, window_end)
            hyps.append(t)
    return(hyps)

def test_hypotheses(hypotheses, obs_times_hash, events):
    """
    For a hypothesis of form (c,e,r,s), test whether c is a potential cause of e related to time window [r, s].
    Also gets *relations*.

    Parameters:
        hypotheses: the hypotheses of form (c,e,r,s) as generated by `generate_hypotheses_for_effects`: a list of tuples
        obs_pair_hash: the dictionary containing all observations by time unit
        obs_times_hash: the dictionary containing all observations by variable
        events: passed from main function, total number of observations

    Returns:
        Dictionary containing all prima facie causes for an effect.
        Key: effect | Val: tuple containing the cause and the start and end of the time window
    """

    logging.log(logging.DEBUG, "Testing hypotheses.")
    # Relations | key: effect | val: tuple (cause, r, s)
    relations = {}

    for hypothesis in pb.progressbar(hypotheses):
        cause, effect, r, s = hypothesis
        c_and_e, c_true, e_true = test_generic_pair_window(cause, effect, r, s, obs_times_hash)

        if is_prima_facie(c_and_e, c_true, e_true, events):
            # add a dictionary key to rel (the effect) containing all causes and their time windows which are prima facie causes
            pf = (cause, r, s)
            if effect in relations:
                relations[effect].append(pf)
            else:
                relations[effect] = [pf]
            
    return(relations)

def test_generic_pair_window(cause, effect, r, s, obs_times_hash):
    """
    Get #(c & e), #c, #e and #events.

    Parameters:
        cause:  the cause variable
        effect: the effect alphabet (list of effects)
        r:      the start of a time window
        s:      the end of a time window
    
    Returns:
        tuple of #(c & e), #c, #e for a given cause, effect and time window.
    """
    # c-mov and e-mov tell the relationship we're testing.
    # get #obs of e-mov within window after c-mov, and #obs of e-mov
    # P(E|C) will be # c->e / # c
    # P(E) will be # e trues / # events

    # First, check formulas for cause and effect
    #utils.check_formula(cause)
    #utils.check_formula(effect)

    c_true_times = utils.get_true_times(obs_times_hash, cause)
    e_true_times = utils.get_true_times(obs_times_hash, effect)
    c_trues = len(c_true_times)
    e_trues = len(e_true_times)
    c_and_e = 0

    for time in c_true_times:
        c_cases = c_true_times[time]
        e_trues_windowed = {key:val for key,val in e_true_times.items() if key >= time + r and key <= time + s}
        e_cases = []
        for list in e_trues_windowed.values():
            e_cases.extend(list)

        for c in c_cases:
            if c in e_cases:
                c_and_e += 1

    return(c_and_e, c_trues, e_trues)

def is_prima_facie(c_and_e, c_true, e_true, events):
    """
    Determines whether c is a prima facie cause of e.

    Parameters:
        c_and_e:    number of times both events were true in a time window
        c_true:     number of times the cause was true in a time window
        e_true:     number of times the effect was true in a window
        events:     the total number of events
    
    Returns:
        boolean
    """
    if c_true == 0:
        return(False)
    
    result = c_and_e / c_true > e_true / events
    return(result)

### -------------------------------
### Define methods - epsilon values
### -------------------------------

def do_all_epsilon_averages(relations, output,  obs_times_hash):
    """
    Get the epsilon values for all relationships

    Parameters:
        relations: dictionary with all prima facie causes for an effect e. Generated by `test_hypotheses()`
        output: the output file to write results to.
        obs_pair_hash: the full dataset containing timestamps & observations.
    
    Returns:
        No return value. Writes a csv with output instead. See `output` parameter.
    """
    message = "effect" if len(relations) == 1 else "effects"
    logging.log(logging.DEBUG, f"Testing significance for {len(relations)} {message}.")

    with open(output, mode='w') as f:
        f.write("cause,effect,w-start,w-end,epsilon")

        for effect in relations:
            logging.log(logging.DEBUG, f"Calculating epsilons for {effect}.")
            for cause, r, s in pb.progressbar(relations[effect]):
                eps = get_epsilon_average(effect, cause, r, s, relations, obs_times_hash)
                # logging.log(logging.DEBUG, f"eps_avg for cause {cause} and effect {effect} is {eps}")

                f.write("\n")
                f.write(f"{cause},{effect},{r},{s},{eps}")
            

def get_epsilon_average(effect, cause, r, s, relations, obs_times_hash):
    """
    Calculates the epsilon value for a given relationship.

    Parameters:
        effect: the variable representing the effect.
        cause: the variable representing the prima facie cause
        r: the start of the time window
        s: the end of the time window
        relations: all prima facie causes to calculate probability differences
    
    Returns:
        Epsilon value for the given relationship.
    """
    other_causes = utils.get_other_causes(effect, cause, relations)

    if len(other_causes) is not 0:
        eps_x = 0
        for x, p, q in other_causes:
            # Sum the epsilon_x for the other causes
            eps_x += utils.calculate_probability_difference(effect, cause, x, p, q, obs_times_hash)
            # logging.log(logging.DEBUG, f"eps_x for cause {cause}, x {x} and effect {effect} is {eps_x}")
        
        # Divide the sum by the amount of other causes: average epsilon to return
        eps_avg = eps_x/len(other_causes)
        return(eps_avg)
    else:
        return None