import random
import time
import numpy as np
import matplotlib.pyplot as plt


class RCAIFinalBenchmark:

    def __init__(
        self,
        vocab_size=2000,
        sequence_length=1000,
        truth_size_range=(100, 150),
        seed=42
    ):

        random.seed(seed)
        np.random.seed(seed)

        self.vocab_size = vocab_size
        self.sequence_length = sequence_length

        self.truth_sets = []

        for _ in range(sequence_length):

            size = random.randint(
                truth_size_range[0],
                truth_size_range[1]
            )

            truth = set(
                random.sample(
                    range(vocab_size),
                    size
                )
            )

            self.truth_sets.append(truth)

    def softmax(self, x):

        x = np.array(x)

        m = np.max(x)

        exp_x = np.exp(x - m)

        return exp_x / exp_x.sum()

    def generate_logits(self):

        return np.random.normal(
            loc=0.0,
            scale=2.0,
            size=self.vocab_size
        )

    def validator_pass(
        self,
        token,
        truth,
        accuracy
    ):

        real = token in truth

        if random.random() < accuracy:
            return real

        return not real

    def build_manifold(
        self,
        truth,
        logic_acc,
        factual_acc,
        semantic_acc
    ):

        manifold = set()

        for token in range(
            self.vocab_size
        ):

            logic_ok = self.validator_pass(
                token,
                truth,
                logic_acc
            )

            factual_ok = self.validator_pass(
                token,
                truth,
                factual_acc
            )

            semantic_ok = self.validator_pass(
                token,
                truth,
                semantic_acc
            )

            if (
                logic_ok
                and factual_ok
                and semantic_ok
            ):
                manifold.add(token)

        return manifold

    def run_baseline(self):

        errors = 0

        start = time.time()

        for truth in self.truth_sets:

            logits = self.generate_logits()

            probs = self.softmax(logits)

            token = np.random.choice(
                self.vocab_size,
                p=probs
            )

            if token not in truth:
                errors += 1

        runtime = (
            time.time() - start
        ) * 1000

        return (
            errors / self.sequence_length,
            runtime
        )

    def run_rcai(
        self,
        validator_accuracy
    ):

        errors = 0

        fp_mass_total = 0

        tp = 0
        fp = 0
        fn = 0

        start = time.time()

        for truth in self.truth_sets:

            manifold = self.build_manifold(
                truth,
                validator_accuracy,
                validator_accuracy,
                validator_accuracy
            )

            if len(manifold) == 0:
                continue

            for token in manifold:

                if token in truth:
                    tp += 1
                else:
                    fp += 1

            for token in truth:

                if token not in manifold:
                    fn += 1

            logits = self.generate_logits()

            projected = np.full(
                self.vocab_size,
                -np.inf
            )

            for token in manifold:
                projected[token] = logits[token]

            probs = self.softmax(
                projected
            )

            fp_mass = 0

            for token in manifold:

                if token not in truth:

                    fp_mass += probs[token]

            fp_mass_total += fp_mass

            sampled = np.random.choice(
                self.vocab_size,
                p=probs
            )

            if sampled not in truth:
                errors += 1

        runtime = (
            time.time() - start
        ) * 1000

        precision = tp / (tp + fp)

        recall = tp / (tp + fn)

        return {

            "hallucination":
                errors /
                self.sequence_length,

            "fp_mass":
                fp_mass_total /
                self.sequence_length,

            "precision":
                precision,

            "recall":
                recall,

            "runtime":
                runtime
        }


def main():

    benchmark = RCAIFinalBenchmark()

    print("=" * 70)
    print("RCAI FINAL VALIDATION BENCHMARK")
    print("=" * 70)

    baseline_error, baseline_time = (
        benchmark.run_baseline()
    )

    print("\nBASELINE")

    print(
        f"Hallucination Rate: "
        f"{baseline_error*100:.2f}%"
    )

    print(
        f"Runtime: "
        f"{baseline_time:.2f} ms"
    )

    accuracies = [
        0.70,
        0.75,
        0.80,
        0.85,
        0.90,
        0.95,
        0.98,
        0.99
    ]

    hallucinations = []
    fp_mass = []

    print("\n")
    print("=" * 70)
    print("RCAI RESULTS")
    print("=" * 70)

    for acc in accuracies:

        result = benchmark.run_rcai(acc)

        hallucinations.append(
            result["hallucination"]
        )

        fp_mass.append(
            result["fp_mass"]
        )

        print(
            f"\nValidator Accuracy: "
            f"{acc:.2f}"
        )

        print(
            f"Hallucination Rate: "
            f"{result['hallucination']*100:.2f}%"
        )

        print(
            f"False Positive Mass: "
            f"{result['fp_mass']:.6f}"
        )

        print(
            f"Precision: "
            f"{result['precision']:.4f}"
        )

        print(
            f"Recall: "
            f"{result['recall']:.4f}"
        )

        print(
            f"Runtime: "
            f"{result['runtime']:.2f} ms"
        )

    plt.figure(
        figsize=(10, 6)
    )

    plt.plot(
        accuracies,
        hallucinations,
        marker="o",
        linewidth=3,
        label="Observed Hallucination Rate"
    )

    plt.plot(
        accuracies,
        fp_mass,
        marker="s",
        linewidth=3,
        label="False Positive Probability Mass"
    )

    plt.xlabel(
        "Validator Accuracy"
    )

    plt.ylabel(
        "Probability"
    )

    plt.title(
        "RCAI Validation Benchmark"
    )

    plt.grid(True)

    plt.legend()

    plt.tight_layout()

    plt.show()


if __name__ == "__main__":
    main()