from httpx import delete
from torch import concat, unsafe_chunk

from cexl.utils import Concept

from ..utils import LocalConceptPredictor
from . import utils
from . import anchor_base
from . import anchor_explanation
import numpy as np
import json
import os
import string
import sys
from io import open
import numpy as np
import time

def id_generator(size=15):
    """Helper function to generate random div ids. This is useful for embedding
    HTML into ipython notebooks."""
    chars = list(string.ascii_uppercase + string.digits)
    return ''.join(np.random.choice(chars, size, replace=True))


def exp_normalize(x):
    b = x.max()
    y = np.exp(x - b)
    return y / y.sum()

class AnchorConcept(object):
    """bla"""

    def __init__(self,  class_names=None, **kwargs):
        """
        Args:
            class_names: list of strings
            use_unk_distribution: if True, the perturbation distribution
                will just replace words randomly with mask_string.
                If False, words will be replaced by similar words using word
                embeddings
            mask_string: String used to mask tokens if use_unk_distribution is True.
        """
        self.class_names = class_names
        self.sample_tm = 0

    def get_sample_fn(self, now_c:np.ndarray, classifier_fn, concepts, onepass=False, use_proba=False) -> tuple[np.ndarray, callable]:
        ""
        true_label = classifier_fn(now_c)[0]
        # words = np.array([x.text for x in processed], dtype='|U80')
        # positions = [x.idx for x in processed]
        def my_sample_fn(present,num_samples,compute_labels=True):        
            data = np.ones((num_samples, len(concepts)))
            # raw = np.zeros((num_samples, len(words)), '|U80')
            # raw[:] = words
            
            b=time.time()
            for i, t in enumerate(concepts):
                if i in present:
                    continue
                n_changed = np.random.binomial(num_samples, .5)
                changed = np.random.choice(num_samples, n_changed,
                                            replace=False)
                # raw[changed, i] = self.mask_string
                data[changed, i] = 0
            # raw_data = [' '.join(x) for x in raw]
            # raw_data = self.inverse_llm(data,concepts)

            self.sample_tm += (time.time() - b)

            labels = []
            if compute_labels:
                b = time.time()
                labels = (classifier_fn(data) == true_label).astype(int)

            labels = np.array(labels)
            # raw_data = raw_data.reshape(-1, 1)
            if not data[:,present].all():
                # print(raw_data)
                print(present)
                print(concepts)
                print(data)
                print(labels)
                # fit(pos)
                raise AssertionError
            # print(pos)
            return data, data, labels
        
        return true_label, my_sample_fn

    def explain_instance(self, predictor:LocalConceptPredictor, concepts:list[Concept], threshold=0.95,
                         delta=0.1, tau=0.15, batch_size=10, onepass=False,
                         use_proba=False, beam_size=4, coverage_samples = 50,
                         **kwargs):
        # print("explain: ", end="")
        # print(rules)
        true_label, sample_fn = self.get_sample_fn(
            np.ones(len(concepts)), predictor.predict_lr, concepts=concepts, onepass=onepass, use_proba=use_proba)
        # print words, true_label
        exp = anchor_base.AnchorBaseBeam.anchor_beam(
            sample_fn, n_features=len(concepts), delta=delta, epsilon=tau, batch_size=batch_size,
            desired_confidence=threshold, stop_on_first=True,
            coverage_samples=coverage_samples,beam_size=beam_size, **kwargs)
        exp['names'] = [f'concept_{x}' for x in exp['feature']]
        # exp['positions'] = [positions[x] for x in exp['feature']]
        # exp['instance'] = text
        exp['prediction'] = true_label
        explanation = anchor_explanation.AnchorExplanation('concept', exp,
                                                           self.as_html)
        return explanation

    def explain_instance_with_distance(self, text, classifier_fn):
        if type(text) == bytes:
            text = text.decode()

    def as_html(self, exp):
        predict_proba = np.zeros(len(self.class_names))
        exp['prediction'] = int(exp['prediction'])
        predict_proba[exp['prediction']] = 1
        predict_proba = list(predict_proba)

        def jsonize(x):
            return json.dumps(x)

        this_dir, _ = os.path.split(__file__)
        bundle = open(os.path.join(this_dir, 'bundle.js'), encoding='utf8').read()
        random_id = 'top_div' + id_generator()

        example_obj = []

        def process_examples(examples, idx):
            idxs = exp['feature'][:idx + 1]
            out_dict = {}
            new_names = {'covered_true': 'coveredTrue', 'covered_false': 'coveredFalse', 'covered': 'covered'}
            for name, new in new_names.items():
                ex = [x[0] for x in examples[name]]
                out = []
                for e in ex:
                    processed = self.nlp(str(e))
                    raw_indexes = [(processed[i].text, processed[i].idx, exp['prediction']) for i in idxs]
                    out.append({'text': e, 'rawIndexes': raw_indexes})
                out_dict[new] = out
            return out_dict

        example_obj = []
        for i, examples in enumerate(exp['examples']):
            example_obj.append(process_examples(examples, i))

        explanation = {'names': exp['names'],
                       'certainties': exp['precision'] if len(exp['precision']) else [exp['all_precision']],
                       'supports': exp['coverage'],
                       'allPrecision': exp['all_precision'],
                       'examples': example_obj}
        processed = self.nlp(exp['instance'])
        raw_indexes = [(processed[i].text, processed[i].idx, exp['prediction'])
                       for i in exp['feature']]
        raw_data = {'text': exp['instance'], 'rawIndexes': raw_indexes}
        jsonize(raw_indexes)

        out = u'''<html>
        <meta http-equiv="content-type" content="text/html; charset=UTF8">
        <head><script>%s </script></head><body>''' % bundle
        out += u'''
        <div id="{random_id}" />
        <script>
            div = d3.select("#{random_id}");
            lime.RenderExplanationFrame(div,{label_names}, {predict_proba},
            {true_class}, {explanation}, {raw_data}, "text", "anchor");
        </script>'''.format(random_id=random_id,
                            label_names=jsonize(self.class_names),
                            predict_proba=jsonize(list(predict_proba)),
                            true_class=jsonize(False),
                            explanation=jsonize(explanation),
                            raw_data=jsonize(raw_data))
        out += u'</body></html>'
        return out

    def show_in_notebook(self, exp, true_class=False, predict_proba_fn=None):
        """Bla"""
        out = self.as_html(exp, true_class, predict_proba_fn)
        from IPython.core.display import display, HTML
        display(HTML(out))
