import torch
import sys
import os

def run_nexus(weights_path):
    sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))

    from src.trainer import load_nexus
    from tokenizers import Tokenizer
    import torch.nn.functional as F

    model, config = load_nexus(weights_path)

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)

    tokenizer_path = os.path.join(os.path.dirname(weights_path), '..', 'data', 'tokenizer.json')
    tokenizer_path = os.path.normpath(tokenizer_path)

    if not os.path.exists(tokenizer_path):
        tokenizer_path = os.path.join(os.path.dirname(__file__), 'data', 'tokenizer.json')

    tokenizer = Tokenizer.from_file(tokenizer_path)

    print("\n{Nexus SmAll v1} Chat Interface")
    print("{Nexus SmAll v1} Type 'exit' to quit, 'clear' to reset conversation")
    print("{Nexus SmAll v1} Type '--temp 0.5' to change temperature")
    print("{Nexus SmAll v1} Type '--help' for all commands\n")

    bos_id = tokenizer.token_to_id("<bos>") if tokenizer.token_to_id("<bos>") is not None else 1
    eos_id = tokenizer.token_to_id("<eos>") if tokenizer.token_to_id("<eos>") is not None else 2

    conversation = [bos_id]

    temperature = 0.2
    top_k = 40
    top_p = 0.9
    max_tokens = 128
    repetition_penalty = 1.2

    while True:
        try:
            user_input = input("You: ").strip()

            if not user_input:
                continue
            if user_input.lower() == 'exit':
                print("Goodbye!")
                break
            elif user_input.lower() == 'clear':
                conversation = [bos_id]
                print("[Conversation reset]")
                continue
            elif user_input.startswith('--'):
                parts = user_input.split()
                if parts[0] == '--temp' and len(parts) >= 2:
                    temperature = float(parts[1])
                    print(f"[temperature={temperature}]")
                    continue
                elif parts[0] == '--help':
                    print("Commands:")
                    print("  --temp <value>    Set temperature (default 0.2)")
                    print("  --topk <value>    Set top_k (default 40)")
                    print("  --topp <value>    Set top_p (default 0.9)")
                    print("  --tokens <value>  Set max new tokens (default 128)")
                    print("  --rep <value>     Set repetition penalty (default 1.2)")
                    print("  clear             Reset conversation")
                    print("  exit              Exit")
                    continue
                elif parts[0] == '--topk' and len(parts) >= 2:
                    top_k = int(parts[1])
                    print(f"[top_k={top_k}]")
                    continue
                elif parts[0] == '--topp' and len(parts) >= 2:
                    top_p = float(parts[1])
                    print(f"[top_p={top_p}]")
                    continue
                elif parts[0] == '--tokens' and len(parts) >= 2:
                    max_tokens = int(parts[1])
                    print(f"[max_tokens={max_tokens}]")
                    continue
                elif parts[0] == '--rep' and len(parts) >= 2:
                    repetition_penalty = float(parts[1])
                    print(f"[repetition_penalty={repetition_penalty}]")
                    continue

            prompt = f"\nUser: {user_input}\nAssistant:"
            prompt_ids = tokenizer.encode(prompt).ids
            input_ids = conversation + prompt_ids

            if len(input_ids) > config.max_seq_len:
                input_ids = input_ids[-config.max_seq_len + 64:]

            input_tensor = torch.tensor([input_ids], dtype=torch.long, device=device)

            generated_ids, full_ids = _generate_with_rep_penalty(
                model, input_tensor, max_new_tokens=max_tokens,
                temperature=temperature, top_k=top_k, top_p=top_p,
                repetition_penalty=repetition_penalty,
                eos_id=eos_id,
            )

            response_ids = full_ids[0, input_tensor.shape[1]:].tolist()
            response_text = tokenizer.decode(response_ids)

            if "<eos>" in response_text:
                response_text = response_text[:response_text.index("<eos>")]
            if "<bos>" in response_text:
                response_text = response_text.replace("<bos>", "")
            if "User:" in response_text:
                response_text = response_text[:response_text.index("User:")]
            if "Assistant:" in response_text:
                response_text = response_text.replace("Assistant:", "")

            response_text = response_text.strip()

            if len(response_text) < 2:
                response_text = "[no response]"

            print(f"Nexus SmAll v1: {response_text}")

            conversation = full_ids[0].tolist()
            if eos_id is not None:
                conversation.append(eos_id)

        except KeyboardInterrupt:
            print("\nGoodbye!")
            break
        except Exception as e:
            print(f"[Error] {e}")
            continue

def _generate_with_rep_penalty(model, input_ids, max_new_tokens, temperature, top_k, top_p, repetition_penalty, eos_id):
    model.eval()

    for _ in range(max_new_tokens):
        seq_len = input_ids.shape[1]
        if seq_len > model.config.max_seq_len:
            input_ids = input_ids[:, -model.config.max_seq_len:]

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

            if repetition_penalty != 1.0:
                for batch_idx in range(logits.shape[0]):
                    for token_idx in range(input_ids.shape[1]):
                        token = input_ids[batch_idx, token_idx].item()
                        if logits[batch_idx, token] < 0:
                            logits[batch_idx, token] *= repetition_penalty
                        else:
                            logits[batch_idx, token] /= repetition_penalty

            logits = logits / temperature

            if top_k > 0:
                top_k_values, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                min_top_k = top_k_values[:, -1].unsqueeze(-1)
                logits = torch.where(logits < min_top_k,
                                     torch.full_like(logits, float('-inf')), logits)

            if top_p > 0 and top_p < 1.0:
                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
                cumulative_probs = torch.cumsum(torch.nn.functional.softmax(sorted_logits, dim=-1), dim=-1)
                sorted_indices_to_remove = cumulative_probs > top_p
                sorted_indices_to_remove[:, 0] = False
                indices_to_remove = torch.zeros_like(logits, dtype=torch.bool)
                indices_to_remove = indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
                logits = torch.where(indices_to_remove,
                                     torch.full_like(logits, float('-inf')), logits)

            probs = torch.nn.functional.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)

        input_ids = torch.cat([input_ids, next_token], dim=-1)

        if eos_id is not None and next_token.item() == eos_id:
            break

    return None, input_ids

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Nexus SmAll v1 Chat")
    parser.add_argument("--weights", type=str, default="weights/nexus_final.pt",
                        help="Path to model weights (.pt file)")
    parser.add_argument("--temp", type=float, default=0.2,
                        help="Temperature (default: 0.2)")
    parser.add_argument("--top_k", type=int, default=40,
                        help="Top-k sampling (default: 40)")
    parser.add_argument("--top_p", type=float, default=0.9,
                        help="Top-p sampling (default: 0.9)")
    parser.add_argument("--max_tokens", type=int, default=128,
                        help="Max new tokens (default: 128)")
    args = parser.parse_args()

    if not os.path.exists(args.weights):
        print(f"[Error] Weights not found: {args.weights}")
        print("Make sure training completed successfully.")
        input("Press Enter to exit...")
        sys.exit(1)

    run_nexus(args.weights)