#!/usr/bin/env python3
"""
NxLook Epistemic AI Inference Library
Author: A.Fatykhov
License: MIT

This library provides inference capabilities for Epistemic AI trained models.
It supports both interactive and non-interactive text generation.

DO NOT MODIFY THIS SCRIPT TO REVERSE ENGINEER TRAINING TECHNIQUES.
This library only performs inference.
"""

import os
import re
import sys
import torch
from torch import nn


class _Model(nn.Module):
    def __init__(self, vocab_size, embedding_dim=512, num_heads=4, num_layers=4, ffn_hidden=256, max_seq_len=128):
        super().__init__()
        self.vocab_size = vocab_size
        self.embedding_dim = embedding_dim
        self.num_heads = num_heads
        self.num_layers = num_layers
        self.ffn_hidden = ffn_hidden
        self.max_seq_len = max_seq_len
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.pos_encoding = nn.Parameter(torch.randn(max_seq_len, embedding_dim) * 0.01)
        self.layers = nn.ModuleList(
            [_TransformerLayer(embedding_dim, num_heads, ffn_hidden) for _ in range(num_layers)])
        self.layer_norm = nn.LayerNorm(embedding_dim)
        self.lm_head = nn.Linear(embedding_dim, vocab_size)

    def forward(self, input_ids):
        embedded = self.embedding(input_ids)
        batch_size, seq_len, _ = embedded.shape
        pos_enc = self.pos_encoding[:seq_len, :].unsqueeze(0)
        embedded = embedded + pos_enc
        for layer in self.layers:
            embedded, _ = layer(embedded)
        embedded = self.layer_norm(embedded)
        logits = self.lm_head(embedded)
        return logits, None


class _TransformerLayer(nn.Module):
    def __init__(self, embedding_dim, num_heads, ffn_hidden):
        super().__init__()
        self.embedding_dim = embedding_dim
        self.attention = nn.MultiheadAttention(embedding_dim, num_heads, batch_first=True)
        self.ffn = nn.Sequential(
            nn.Linear(embedding_dim, ffn_hidden),
            nn.ReLU(),
            nn.Linear(ffn_hidden, embedding_dim)
        )
        self.norm1 = nn.LayerNorm(embedding_dim)
        self.norm2 = nn.LayerNorm(embedding_dim)
        self.dropout = nn.Dropout(0.1)

    def forward(self, x):
        attn_output, _ = self.attention(x, x, x)
        x = self.norm1(x + self.dropout(attn_output))
        ffn_output = self.ffn(x)
        x = self.norm2(x + self.dropout(ffn_output))
        return x, None


def _tokenize(text):
    if not isinstance(text, str):
        text = ""
    tokens = re.sub(r'\s+', ' ', text.lower()).strip().split()
    unique_tokens = []
    seen = set()
    for token in tokens:
        if token not in seen:
            unique_tokens.append(token)
            seen.add(token)
    vocab = {'<pad>': 0, '<unk>': 1, '<start>': 2, '<end>': 3}
    for token in unique_tokens:
        if token not in vocab:
            vocab[token] = len(vocab)
    id_to_token = {i: t for t, i in vocab.items()}
    return vocab, id_to_token


def _encode(text, stoi):
    tokens = text.lower().split()
    return [stoi.get(token, 1) for token in tokens]


def _decode(ids, itos):
    tokens = [itos.get(i, '<unk>') for i in ids]
    return ' '.join(tokens).strip()


def _generate_text(model, stoi, itos, prompt, max_new_tokens=100, temperature=0.7, top=40, cut_off=10):
    device = torch.device(
        "mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    model.eval()

    required_tokens = {'<pad>', '<unk>', '<start>', '<end>'}
    if not all(t in stoi for t in required_tokens):
        raise RuntimeError("Tokenizer is corrupted. Missing required tokens.")

    tokens = prompt.lower().split()
    input_ids = [stoi.get(t, 1) for t in tokens]
    if '<start>' in stoi:
        input_ids = [stoi['<start>']] + input_ids

    pad_id = stoi.get('<pad>', 0)
    max_seq_len = model.max_seq_len
    if len(input_ids) < max_seq_len:
        input_ids = [pad_id] * (max_seq_len - len(input_ids)) + input_ids

    generated = input_ids.copy()
    past_tokens = []
    seen_sequences = set()
    end_count = 0

    try:
        for _ in range(max_new_tokens):
            # Truncate if too long
            start_idx = max(0, len(generated) - max_seq_len)
            input_tensor = torch.tensor([generated[start_idx:start_idx + max_seq_len]], dtype=torch.long).to(device)

            with torch.no_grad():
                logits, _ = model(input_tensor)
                next_token_logits = logits[0, -1, :]

            next_token_logits = next_token_logits / temperature

            for past_id in past_tokens[-5:]:
                if 0 <= past_id < len(next_token_logits):
                    next_token_logits[past_id] /= 1.5

            k = min(top, next_token_logits.size(-1))
            topk_probs, topk_indices = torch.topk(torch.softmax(next_token_logits, dim=-1), k)

            cumulative_probs = torch.cumsum(topk_probs, dim=0)
            cutoff_idx = (cumulative_probs >= 0.85).nonzero(as_tuple=False)
            if len(cutoff_idx) > 0:
                cutoff = cutoff_idx[0].item()
            else:
                cutoff = len(topk_probs) - 1

            topk_probs = topk_probs[:cutoff + 1]
            topk_indices = topk_indices[:cutoff + 1]

            aw = topk_probs.clone()
            for i, idx in enumerate(topk_indices):
                candidate_seq = generated[-8:] + [idx.item()]
                seq_str = ' '.join([itos.get(tid, '<unk>') for tid in candidate_seq])
                if seq_str in seen_sequences:
                    aw[i] *= 0.5
                else:
                    seen_sequences.add(seq_str)

            topk_probs = torch.softmax(aw, dim=0)
            next_id = torch.multinomial(topk_probs, num_samples=1).item()
            next_id = topk_indices[next_id].item()

            generated.append(next_id)
            past_tokens.append(next_id)
            if len(past_tokens) > 5:
                past_tokens.pop(0)

            token = itos.get(next_id, '<unk>')
            if token in {'.', '!', '?'}:
                end_count += 1
                if end_count >= 2:
                    break

            if next_id == stoi.get('<end>', -1) or len(generated) >= max_seq_len + max_new_tokens:
                break

    except Exception as e:
        return f"Generation failed: {str(e)}"

    # Decode and clean output
    text = _decode(generated, itos)
    words = [w for w in text.split() if w not in {'<start>', '<end>', '<pad>', '<unk>'}]
    cleaned = ' '.join(words)

    if len(cleaned.strip()) < cut_off:
        return "I cannot confidently predict."

    # Deduplicate trigrams
    final_words = []
    seen_trigrams = set()
    for i in range(len(words)):
        if i >= 2:
            trigram = ' '.join(words[i - 2:i + 1])
            if trigram in seen_trigrams:
                continue
            seen_trigrams.add(trigram)
        final_words.append(words[i])

    return ' '.join(final_words)


def _load_model(model_path):
    """Private: Load model and tokenizer from .pth file."""
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"Model file not found: {model_path}")

    checkpoint = torch.load(model_path, map_location='cpu')

    if 'hyperparameters' in checkpoint:
        hp = checkpoint['hyperparameters']
    elif isinstance(checkpoint.get('model_state_dict'), dict):
        hp = {
            'embedding_dim': checkpoint['model_state_dict']['embedding.weight'].shape[1],
            'num_heads': 4,
            'num_layers': len(
                [k for k in checkpoint['model_state_dict'].keys() if '.layers.' in k and 'ffn.0.weight' in k]),
            'ffn_hidden': checkpoint['model_state_dict']['layers.0.ffn.0.weight'].shape[0],
            'max_seq_len': checkpoint['model_state_dict']['pos_encoding'].shape[0],
            'vocab_size': checkpoint['vocab_size']
        }
    else:
        raise RuntimeError("Checkpoint does not contain valid model configuration.")

    vocab_size = checkpoint['vocab_size']
    stoi = checkpoint['stoi']
    itos = checkpoint['itos']

    model = _Model(
        vocab_size=vocab_size,
        embedding_dim=hp['embedding_dim'],
        num_heads=hp.get('num_heads', 4),
        num_layers=hp.get('num_layers', 8),
        ffn_hidden=hp['ffn_hidden'],
        max_seq_len=hp['max_seq_len']
    )
    model.load_state_dict(checkpoint['model_state_dict'])

    print(f"✅ Model loaded from {model_path}")
    print(f"   Architecture: vocab={vocab_size}, emb_dim={hp['embedding_dim']}, "
          f"ffn_hidden={hp['ffn_hidden']}, layers={hp['num_layers']}, max_seq_len={hp['max_seq_len']}")

    return model, stoi, itos

def interactive_inference(model_path, temperature=0.7):
    """
    Start an interactive chat session with the model.

    Args:
        model_path (str): Path to the .pth model file.
        temperature (float): Sampling temperature. Default: 0.7.

    Example:
        >>> interactive_inference("model.pth")
        👤 You: Hello
        🤖 Model:  Hi there! How can I help you today?
    """
    print(f"\n🤖 Q&A INFERENCE MODE (Loading model from: {model_path})")
    try:
        model, stoi, itos = _load_model(model_path)
    except Exception as e:
        print(f"❌ Failed to load model: {e}")
        sys.exit(1)
    print("✅ Model loaded. Ready for interaction.")
    print("💡 Type 'quit' or 'exit' to leave.")

    while True:
        print("\n" + "=" * 60)
        prompt = input("👤 You: ").strip()
        if not prompt:
            continue
        if prompt.lower() in ['quit', 'exit']:
            print("👋 Goodbye!")
            break
        print("🤖 Model:")
        response = _generate_text(model, stoi, itos, prompt, 100, temperature)
        print(f"  {response}\n")


def non_interactive_inference(model_path, prompt, temperature=0.7):
    """
    Perform a single non-interactive inference.

    Args:
        model_path (str): Path to the .pth model file.
        prompt (str): Input text to generate from.
        temperature (float): Sampling temperature. Default: 0.7.

    Returns:
        str: Generated text response.

    Example:
        >>> response = non_interactive_inference("model.pth", "What is AI?")
        >>> print(response)
    """
    if not isinstance(prompt, str):
        raise TypeError("prompt must be a string")

    print(f"\n🤖 NON-INTERACTIVE INFERENCE MODE")
    print(f"  Model: {model_path}")
    print(f"  Prompt: '{prompt}'")

    try:
        model, stoi, itos = _load_model(model_path)
    except Exception as e:
        print(f"❌ Failed to load model: {e}")
        sys.exit(1)

    response = _generate_text(model, stoi, itos, prompt, 100, temperature)
    print(f"\n🪄 Generated Response:")
    print(response)

    return response

# ======================
# PUBLIC API FUNCTIONS (these are what users import)
# ======================

def inference(model_path, prompt=None, interactive=False, temperature=0.7):
    """
    Perform LLM inference using a trained Epistemic AI model.

    Args:
        model_path (str): Path to the .pth model file.
        prompt (str, optional): Prompt for non-interactive inference. Ignored if interactive=True.
        interactive (bool): If True, enter interactive chat mode.
        temperature (float): Sampling temperature (0.1 to 1.0). Default: 0.7.

    Returns:
        str: Generated text (in non-interactive mode), or None in interactive mode.

    Example:
        >>> response = inference("model.pth", prompt="What is the human brain?")
        >>> print(response)
    """
    if not isinstance(temperature, float) or not (0.1 <= temperature <= 1.0):
        raise ValueError("temperature must be a float between 0.1 and 1.0")

    if interactive:
        interactive_inference(model_path, temperature)
        return None  # Interactive mode prints directly
    elif prompt is not None:
        return non_interactive_inference(model_path, prompt, temperature)
    else:
        raise ValueError("Either 'prompt' must be provided or 'interactive=True'")


# ======================
# CLI Entry Point (for command-line use)
# ======================

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="NxLook Epistemic AI Inference Tool")
    parser.add_argument("--model", required=True, help="Path to .pth model file")
    parser.add_argument("--prompt", type=str, help="Prompt for inference (non-interactive)")
    parser.add_argument("--interactive", action="store_true", help="Start interactive chat mode")
    parser.add_argument("--temperature", type=float, default=0.7,
                        help="Sampling temperature (default: 0.7)")

    args = parser.parse_args()

    if not os.path.exists(args.model):
        print(f"❌ Model file not found: {args.model}")
        sys.exit(1)

    if args.interactive:
        interactive_inference(args.model, args.temperature)
    elif args.prompt is not None:
        non_interactive_inference(args.model, args.prompt, args.temperature)
    else:
        print("❌ Please provide either --prompt or --interactive.")
        sys.exit(1)
