import numpy as np
import pandas as pd

from pyqubo import Array, Placeholder, Constraint

def prob(dataset:pd.DataFrame, max_bins=10):
    """Joint probability distribution P(X) for the given data.
    
    Parameters:
    dataset (DataFrame): Input DataFrame
    max_bins (int): Maximal number of bins to discretize the sample

    Returns:
    list: Join probability distribution
    """

    # bin by the number of different values per feature
    _, num_columns = dataset.shape
    bins = [min(len(np.unique(dataset[:, ci])), max_bins) for ci in range(num_columns)]

    freq, _ = np.histogramdd(dataset, bins)
    joint_prob = freq / np.sum(freq)
    return joint_prob

def shannon_entropy(joint_prob) -> int:
    """Shannon entropy H(X) is the sum of P(X)log(P(X)) for probabilty distribution P(X)."""
    flatten_probs = joint_prob.flatten()
    return -sum(pi*np.log2(pi) for pi in flatten_probs if pi)

def conditional_shannon_entropy(p, *conditional_indices):
    """Shannon entropy of P(X) conditional on variable j"""

    axis = tuple(i for i in np.arange(len(p.shape)) if i not in conditional_indices)
    return shannon_entropy(p) - shannon_entropy(np.sum(p, axis=axis))

def mutual_information(p, j):
    """Mutual information between all variables and variable j"""
    return shannon_entropy(np.sum(p, axis=j)) - conditional_shannon_entropy(p, j)

def conditional_mutual_information(p, j, *conditional_indices):
    """Mutual information between variables X and variable Y conditional on variable Z."""

    marginal_conditional_indices = [i-1 if i > j else i for i in conditional_indices]
    return (conditional_shannon_entropy(np.sum(p, axis=j), *marginal_conditional_indices)
            - conditional_shannon_entropy(p, j, *conditional_indices))


def build_qubo(input_ds:pd.DataFrame, target:str, max_feat:int = None):
    """
    Takes the input data and creates a pyqubo model so that Ising
    coefficients can be extracted. Afterwards this can be taken to
    Kipu's DQA or DCQO routines.
    """

    col_list = input_ds.columns
    feat_list = [ col for col in col_list if col != target]
    num_assets = len(feat_list)
    x = Array.create('x', shape=num_assets, vartype='BINARY')

    # Profit generated by each asset individually
    H_importance = 0.0
    for i, column in enumerate(feat_list):
        if column != target:
            minf = mutual_information(
                prob(input_ds[[target, column]].values), 1)
            H_importance += Constraint(
                minf * x[i], label='profit({})'.format(i)
            )

    # Risk obtained from the covariance matrix
    H_redundancy = 0.0
    shift = Placeholder('shift')
    for i, field0 in enumerate(feat_list):
        for j, field1 in enumerate(feat_list):
            if field0 != field1:
                cmi_01 = conditional_mutual_information(
                    prob(input_ds[[target, field0, field1]].values), 1, 2)
                H_redundancy += Constraint(
                    (cmi_01-shift) * x[i] * x[j], label='risk({}, {})'.format(i, j)
                )

    # Constraint (budget)
    H_penalty = 0.0
    for i in range(num_assets):
        H_penalty += Constraint(x[i], label='slot({})'.format(i))

    # Build model
    lamda = Placeholder('lambda')
    H =  H_redundancy - H_importance + lamda * (H_penalty - max_feat)**2

    return H.compile()

def get_ising_coeffs(dataframe: pd.DataFrame, target:str, max_feat:int = None, reg_lambda: float = 5.0, shift: float = 0.0):

    # Build a QUBO out of provided dataframe
    model = build_qubo(dataframe, target, max_feat)

    # Get coefficients
    feed_dict = {'lambda': reg_lambda, 'shift' : shift}
    ising_coeffs = model.to_ising(feed_dict=feed_dict)

    # Format the coefs
    h = []
    jp = {}
    linear = ising_coeffs[0]
    quadratic = ising_coeffs[1]
    for i in range(len(linear)):
        h += [linear[f'x[{i}]']]
        for j in range(len(linear)):
            if (f'x[{i}]',f'x[{j}]') in quadratic:
                jp[(f"{i}-{j}")] = quadratic[(f'x[{i}]',f'x[{j}]')]

    return h, jp
