import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import os
import time
import threading
import glob
import re
import json
import numpy as np
import wikipediaapi

# ============================================================
#  5x2T — Word Level Model with Disk Offloaded Optimizer
#  Uses numpy for disk saves — much more memory efficient
# ============================================================

# ---------------- CONFIG ----------------
MAX_TRAIN_MIN   = 60
BATCH_SIZE      = 2
SEQ_LENGTH      = 64
EMBED_SIZE      = 192
HIDDEN_SIZE     = 384
NUM_LAYERS      = 1
DROPOUT         = 0.2
LEARNING_RATE   = 0.01
GRAD_CLIP       = 1.0
AUTOSAVE_MIN    = 5
MAX_VOCAB       = 995600
TEMPERATURE     = 0.8
RESPONSE_LENGTH = 40
MAX_LENGTH      = 200
UNK_TOKEN       = "<UNK>"
PAD_TOKEN       = "<PAD>"

BASE_DIR      = r"C:\Users\Eclipsed\Downloads\5x2T"
DATASET_DIR   = os.path.join(BASE_DIR, "datasets")
MODEL_DIR     = os.path.join(BASE_DIR, "5x2T-2")
MODEL_PATH    = os.path.join(MODEL_DIR, "model.pth")
VOCAB_PATH    = os.path.join(MODEL_DIR, "vocab.json")
OFFLOAD_DIR   = os.path.join(MODEL_DIR, "offload")

DATASET_FOLDERS = [
    os.path.join(DATASET_DIR, "chat_dataset"),
    os.path.join(DATASET_DIR, "python_data"),
    os.path.join(DATASET_DIR, "lua_dataset"),
    os.path.join(DATASET_DIR, "dic_dataset"),
    os.path.join(DATASET_DIR, "Wiki_dataset"),
    r"E:\5x2T",
]

DEVICE = torch.device("cpu")


# ---------------- DISK OFFLOAD OPTIMIZER ----------------
class DiskOffloadSGD:
    """
    SGD optimizer that stores momentum buffers on disk as numpy files.
    Large buffers are processed in chunks to avoid RAM spikes.
    """
    CHUNK = 4_000_000  # process 4 million elements at a time

    def __init__(self, params, lr=0.01, momentum=0.9, offload_dir=OFFLOAD_DIR):
        self.params      = list(params)
        self.lr          = lr
        self.momentum    = momentum
        self.offload_dir = offload_dir
        os.makedirs(offload_dir, exist_ok=True)

        print(f"  Initialising {len(self.params)} momentum buffers on disk...")
        for i, p in enumerate(self.params):
            path = os.path.join(offload_dir, f"m_{i}.npy")
            if not os.path.exists(path):
                # Save in chunks to avoid allocating the full array at once
                shape = p.data.shape
                total = p.data.numel()
                flat  = np.zeros(total, dtype=np.float32)
                np.save(path, flat.reshape(shape))
                del flat
        print(f"  Momentum buffers ready in {offload_dir}\n")

    def zero_grad(self):
        for p in self.params:
            if p.grad is not None:
                p.grad.detach_()
                p.grad.zero_()

    def step(self):
        for i, p in enumerate(self.params):
            if p.grad is None:
                continue

            path  = os.path.join(self.offload_dir, f"m_{i}.npy")
            shape = p.data.shape
            total = p.data.numel()

            # Use memory mapped file — only the chunk we touch is in RAM
            buf_mm  = np.load(path, mmap_mode="r+")
            buf_flat = buf_mm.reshape(-1)
            grad_flat = p.grad.data.reshape(-1).numpy()
            data_flat = p.data.reshape(-1).numpy()

            # Process in chunks so RAM never spikes
            for start in range(0, total, self.CHUNK):
                end = min(start + self.CHUNK, total)
                buf_flat[start:end]  = (self.momentum * buf_flat[start:end]
                                        + grad_flat[start:end])
                data_flat[start:end] -= self.lr * buf_flat[start:end]

            # Write updated data back to param tensor
            p.data.copy_(torch.from_numpy(data_flat.reshape(shape)))

            # Flush mmap and free
            buf_mm.flush()
            del buf_mm, buf_flat, grad_flat, data_flat

    def state_dict(self):
        return {"lr": self.lr, "momentum": self.momentum}

    def load_state_dict(self, state):
        self.lr       = state.get("lr", self.lr)
        self.momentum = state.get("momentum", self.momentum)


# ---------------- DATASET DISCOVERY ----------------
def find_all_txt_files(folders):
    all_files = []
    for folder in folders:
        if os.path.exists(folder):
            found = glob.glob(os.path.join(folder, "**", "*.txt"), recursive=True)
            all_files.extend(found)
            print(f"  [{os.path.basename(folder)}] -> {len(found)} file(s)")
        else:
            print(f"  [SKIP] Not found: {folder}")
    if not all_files:
        raise FileNotFoundError("No .txt files found. Check your dataset paths.")
    print(f"\n  Total files: {len(all_files)}\n")
    return all_files


# ---------------- TOKENISER ----------------
def tokenise(text):
    return re.findall(r"\b\w+\b|[\"'.,!?;:\-\n]", text.lower())


def build_vocab(files, max_vocab=MAX_VOCAB):
    print("  Building vocabulary...")
    freq = {}
    total_tokens = 0
    for f in files:
        try:
            with open(f, "r", encoding="utf-8", errors="ignore") as file:
                tokens = tokenise(file.read())
                for t in tokens:
                    freq[t] = freq.get(t, 0) + 1
                total_tokens += len(tokens)
        except Exception as e:
            print(f"  [WARNING] Could not read {f}: {e}")

    sorted_vocab = sorted(freq.items(), key=lambda x: x[1], reverse=True)
    vocab_words  = [PAD_TOKEN, UNK_TOKEN] + [w for w, _ in sorted_vocab[:max_vocab - 2]]
    word2idx     = {w: i for i, w in enumerate(vocab_words)}
    idx2word     = {i: w for i, w in enumerate(vocab_words)}

    print(f"  Total tokens    : {total_tokens:,}")
    print(f"  Unique words    : {len(freq):,}")
    print(f"  Vocab size      : {len(vocab_words):,}\n")

    return vocab_words, word2idx, idx2word


def save_vocab(vocab_words, path):
    with open(path, "w", encoding="utf-8") as f:
        json.dump(vocab_words, f)


def load_vocab(path):
    with open(path, "r", encoding="utf-8") as f:
        vocab_words = json.load(f)
    word2idx = {w: i for i, w in enumerate(vocab_words)}
    idx2word = {i: w for i, w in enumerate(vocab_words)}
    return vocab_words, word2idx, idx2word


# ---------------- DATASET ----------------
class WordDataset(Dataset):
    def __init__(self, files, word2idx):
        self.data = []
        unk_idx   = word2idx.get(UNK_TOKEN, 1)
        for f in files:
            try:
                with open(f, "r", encoding="utf-8", errors="ignore") as file:
                    tokens     = tokenise(file.read())
                    self.data += [word2idx.get(t, unk_idx) for t in tokens]
            except Exception as e:
                print(f"  [WARNING] Could not read {f}: {e}")

        if not self.data:
            raise ValueError("Dataset is empty after tokenisation.")
        print(f"  Dataset tokens: {len(self.data):,}\n")

    def __len__(self):
        return len(self.data) - SEQ_LENGTH

    def __getitem__(self, idx):
        x = torch.tensor(self.data[idx:idx + SEQ_LENGTH],         dtype=torch.long)
        y = torch.tensor(self.data[idx + 1:idx + SEQ_LENGTH + 1], dtype=torch.long)
        return x, y


# ---------------- MODEL ----------------
class Model(nn.Module):
    def __init__(self, vocab_size):
        super().__init__()
        self.embed   = nn.Embedding(vocab_size, EMBED_SIZE, padding_idx=0)
        self.dropout = nn.Dropout(DROPOUT)
        self.lstm    = nn.LSTM(
            EMBED_SIZE, HIDDEN_SIZE,
            num_layers=NUM_LAYERS,
            batch_first=True,
            dropout=0
        )
        self.norm = nn.LayerNorm(HIDDEN_SIZE)
        self.fc   = nn.Linear(HIDDEN_SIZE, vocab_size)

    def forward(self, x, hc=None):
        x = self.dropout(self.embed(x))
        x, hc = self.lstm(x, hc)
        x = self.norm(x)
        x = self.fc(x)
        return x, hc


# ---------------- SETUP ----------------
def setup_dirs():
    os.makedirs(MODEL_DIR,   exist_ok=True)
    os.makedirs(OFFLOAD_DIR, exist_ok=True)
    os.makedirs(os.path.join(MODEL_DIR, "questions"), exist_ok=True)


# ---------------- GENERATE ----------------
def generate(model, word2idx, idx2word, seed_text, length=RESPONSE_LENGTH, temperature=TEMPERATURE):
    model.eval()
    tokens  = tokenise(seed_text)
    unk_idx = word2idx.get(UNK_TOKEN, 1)
    indices = [word2idx.get(t, unk_idx) for t in tokens]
    hc      = None

    with torch.no_grad():
        for _ in range(min(length, MAX_LENGTH)):
            x        = torch.tensor([indices[-SEQ_LENGTH:]], dtype=torch.long)
            out, hc  = model(x, hc)
            logits   = out[0, -1] / temperature
            probs    = torch.softmax(logits, dim=0)
            next_idx = torch.multinomial(probs, 1).item()
            indices.append(next_idx)

    generated = indices[len(tokens):]
    words     = [idx2word.get(i, UNK_TOKEN) for i in generated]
    return " ".join(words)


def format_response(text):
    text = re.sub(r' ([.,!?;:])', r'\1', text)
    text = re.sub(r'\n ', '\n', text)
    if text:
        text = text[0].upper() + text[1:]
    return text


# ---------------- WIKIPEDIA ----------------
wiki_api = wikipediaapi.Wikipedia(
    language='en',
    extract_format=wikipediaapi.ExtractFormat.WIKI,
    user_agent="5x2T-AI/1.0"
)

def search_wikipedia(query):
    try:
        search_term = query.lower()
        for prefix in ["what is ", "what are ", "who is ", "who was ",
                        "tell me about ", "explain ", "define ",
                        "what was ", "how does ", "how do "]:
            search_term = search_term.replace(prefix, "")
        search_term = search_term.replace("?", "").strip()
        page = wiki_api.page(search_term)
        if page.exists():
            return f"[Wikipedia: {page.title}]\n{page.summary[:600]}"
        return None
    except Exception as e:
        print(f"  [WARNING] Wikipedia lookup failed: {e}")
        return None


def should_search_wiki(text):
    triggers = [
        "what is", "what are", "who is", "who was",
        "tell me about", "explain", "define", "what was",
        "how does", "how do"
    ]
    return any(text.lower().strip().startswith(t) for t in triggers)


# ---------------- TRAINING ----------------
def train():
    setup_dirs()
    print("=" * 55)
    print("  5x2T — Word Level Training (Disk Offload)")
    print(f"  Device      : {DEVICE}")
    print(f"  Offload dir : {OFFLOAD_DIR}")
    print(f"  Target      : {MAX_TRAIN_MIN} minutes")
    print("=" * 55 + "\n")

    print("Scanning dataset folders...")
    files = find_all_txt_files(DATASET_FOLDERS)

    if os.path.exists(VOCAB_PATH):
        print("  Found existing vocab — loading...")
        vocab_words, word2idx, idx2word = load_vocab(VOCAB_PATH)
        print(f"  Vocab size: {len(vocab_words):,}\n")
    else:
        vocab_words, word2idx, idx2word = build_vocab(files)
        save_vocab(vocab_words, VOCAB_PATH)
        print(f"  Vocab saved to {VOCAB_PATH}\n")

    print("Loading dataset...")
    dataset = WordDataset(files, word2idx)
    loader  = DataLoader(
        dataset, batch_size=BATCH_SIZE,
        shuffle=True, num_workers=0
    )

    vocab_size = len(vocab_words)
    model      = Model(vocab_size)
    criterion  = nn.CrossEntropyLoss(ignore_index=0)
    optimizer  = DiskOffloadSGD(
        model.parameters(),
        lr=LEARNING_RATE,
        momentum=0.9,
        offload_dir=OFFLOAD_DIR
    )

    param_count = sum(p.numel() for p in model.parameters())
    print(f"  Model parameters : {param_count:,}")
    print(f"  Vocab size       : {vocab_size:,}")
    print(f"  Optimizer        : DiskOffloadSGD (numpy on disk)")
    print(f"  Offload folder   : {OFFLOAD_DIR}\n")

    if os.path.exists(MODEL_PATH + ".npz"):
        load_path = MODEL_PATH + ".npz"
    elif os.path.exists(MODEL_PATH):
        load_path = MODEL_PATH
    else:
        load_path = None

    if load_path:
        try:
            if load_path.endswith(".npz"):
                raw         = np.load(load_path)
                checkpoint  = {k: torch.from_numpy(raw[k]) for k in raw.files}
            else:
                checkpoint  = torch.load(load_path, map_location="cpu")
            model_state = model.state_dict()
            loaded = 0
            for k in checkpoint.keys():
                if k in model_state and checkpoint[k].shape == model_state[k].shape:
                    model_state[k] = checkpoint[k]
                    loaded += 1
            model.load_state_dict(model_state)
            print(f"  Resumed from checkpoint ({loaded} layers matched)\n")
        except Exception as e:
            print(f"  Could not load checkpoint: {e} — starting fresh\n")

    print("-" * 55)
    print("  Training started...\n")

    start_time    = time.time()
    epoch         = 0
    best_loss     = float("inf")
    total_tokens  = 0
    last_autosave = 0
    epoch_loss    = 0
    batches       = 0
    loss          = None

    def print_progress():
        while True:
            elapsed_sec = time.time() - start_time
            elapsed_min = elapsed_sec / 60
            speed = total_tokens / (elapsed_sec + 1e-5)
            avg_loss = epoch_loss / max(batches, 1) if batches > 0 else 0
            mins = int(elapsed_sec // 60)
            secs = int(elapsed_sec % 60)
            current_loss = loss.item() if loss is not None else 0.0
            print(
                f"  Epoch {epoch+1:>3} | "
                f"Batch {batches:>5} | "
                f"Loss: {current_loss:.4f} | "
                f"Avg: {avg_loss:.4f} | "
                f"Speed: {speed:.0f} tok/s | "
                f"Time: {mins:02d}:{secs:02d}/{MAX_TRAIN_MIN:02d}:00",
                end="\r"
            )
            time.sleep(1)

    # Start progress printing thread
    progress_thread = threading.Thread(target=print_progress, daemon=True)
    progress_thread.start()

    while (time.time() - start_time) / 60 < MAX_TRAIN_MIN:
        epoch_loss = 0
        batches    = 0

        for x, y in loader:
            optimizer.zero_grad()
            out, _ = model(x)
            out    = out.view(-1, vocab_size)
            y      = y.view(-1)
            loss   = criterion(out, y)
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)
            optimizer.step()

            epoch_loss   += loss.item()
            batches      += 1
            total_tokens += x.numel()

            elapsed_min = (time.time() - start_time) / 60
            mins = int(elapsed_min)
            secs = int((elapsed_min - mins) * 60)
            if elapsed_min - last_autosave >= AUTOSAVE_MIN:
                try:
                    torch.save(model.state_dict(), MODEL_PATH)
                    last_autosave = elapsed_min
                    print(f"\n  [Autosave] {mins:02d}:{secs:02d} -> {MODEL_PATH}")
                except MemoryError:
                    try:
                        print(f"\n  [Autosave] RAM full - saving directly to disk...")
                        tmp_path = MODEL_PATH + ".tmp"
                        with open(tmp_path, "wb") as f:
                            state = {k: v.numpy() for k, v in model.state_dict().items()}
                            np.savez_compressed(f, **state)
                        os.replace(tmp_path, MODEL_PATH + ".npz")
                        last_autosave = elapsed_min
                        print(f"\n  [Autosave] {mins:02d}:{secs:02d} -> {MODEL_PATH}.npz")
                    except OSError:
                        print(f"\n  [5xSc-404] Low storage or memory - autosave skipped")
                    except Exception as e:
                        print(f"\n  [5xSc-9512] Unknown autosave error: {e}")
                except OSError:
                    print(f"\n  [5xSc-404] Low storage or memory - autosave skipped")
                except KeyboardInterrupt:
                    print(f"\n  [5xSc-80082] Training stopped early - saving...")
                    try:
                        torch.save(model.state_dict(), MODEL_PATH)
                    except Exception:
                        state = {k: v.numpy() for k, v in model.state_dict().items()}
                        np.savez_compressed(MODEL_PATH + ".npz", **state)
                    print(f"  Model saved. Exiting.")
                    raise
                except Exception as e:
                    err = str(e).lower()
                    if "corrupt" in err or "invalid" in err:
                        print(f"\n  [5xSc-312] Corruption detected: {e}")
                    elif "allocat" in err or "memory" in err:
                        print(f"\n  [5xSc-500] Memory allocation failed: {e}")
                    else:
                        print(f"\n  [5xSc-9512] Unknown error: {e}")
            if elapsed_min >= MAX_TRAIN_MIN:
                break

        print()
        epoch    += 1
        avg_loss  = epoch_loss / max(batches, 1)

        if avg_loss < best_loss:
            best_loss = avg_loss
            try:
                torch.save(model.state_dict(), MODEL_PATH)
            except MemoryError:
                print(f"  [Save] RAM full — saving directly to disk...")
                tmp_path = MODEL_PATH + ".tmp"
                with open(tmp_path, "wb") as f:
                    state = {k: v.numpy() for k, v in model.state_dict().items()}
                    np.savez_compressed(f, **state)
                os.replace(tmp_path, MODEL_PATH + ".npz")
            print(f"  [Saved] Best loss: {best_loss:.4f}\n")

        if (time.time() - start_time) / 60 >= MAX_TRAIN_MIN:
            break

    print("-" * 55)
    print(f"  Done! Epochs: {epoch} | Best loss: {best_loss:.4f}")
    print(f"  Model saved to: {MODEL_PATH}\n")

    print("  Sample generation:")
    seed   = '"what is marxism"\n"'
    sample = generate(model, word2idx, idx2word, seed_text=seed, length=40)
    print(f"  {format_response(sample)}\n")

    return model, word2idx, idx2word


# ---------------- CHAT ----------------
def chat(model=None, word2idx=None, idx2word=None):
    print("=" * 55)
    print("  5x2T — Chat")
    print("  Commands:")
    print("    quit       — exit")
    print("    temp X     — temperature e.g. temp 0.7")
    print("    length X   — response length e.g. length 60")
    print("    maxlen X   — max length cap e.g. maxlen 300")
    print("    wiki X     — force Wikipedia lookup e.g. wiki Python")
    print("=" * 55 + "\n")

    if model is None:
        if not os.path.exists(VOCAB_PATH):
            print("[ERROR] No vocab found. Run training first.")
            return
        if not os.path.exists(MODEL_PATH):
            print("[ERROR] No model found. Run training first.")
            return
        vocab_words, word2idx, idx2word = load_vocab(VOCAB_PATH)
        model = Model(len(vocab_words))
        model.load_state_dict(torch.load(MODEL_PATH, map_location="cpu"))
        model.eval()
        param_count = sum(p.numel() for p in model.parameters())
        print(f"  Vocab size       : {len(vocab_words):,}")
        print(f"  Model parameters : {param_count:,}")
        print(f"  Device           : {DEVICE}\n")

    temperature     = TEMPERATURE
    response_length = RESPONSE_LENGTH
    max_length      = MAX_LENGTH

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

        if not user_input:
            continue
        if user_input.lower() in ("quit", "exit", "q"):
            print("Goodbye.")
            break
        if user_input.lower().startswith("temp "):
            try:
                temperature = float(user_input.split()[1])
                print(f"  Temperature -> {temperature}\n")
            except:
                print("  Usage: temp 0.8\n")
            continue
        if user_input.lower().startswith("length "):
            try:
                response_length = int(user_input.split()[1])
                print(f"  Length -> {response_length}\n")
            except:
                print("  Usage: length 50\n")
            continue
        if user_input.lower().startswith("maxlen "):
            try:
                max_length = int(user_input.split()[1])
                print(f"  Max length -> {max_length}\n")
            except:
                print("  Usage: maxlen 300\n")
            continue

        if user_input.lower().startswith("wiki "):
            query  = user_input[5:].strip()
            result = search_wikipedia(query)
            reply  = result if result else f"No Wikipedia page found for '{query}'"
            print(f"5x2T: {reply}\n")
            continue

        wiki_result = None
        if should_search_wiki(user_input):
            wiki_result = search_wikipedia(user_input)

        if wiki_result:
            print(f"5x2T: {wiki_result}\n")
        else:
            seed  = f'"{user_input.lower()}"\n"'
            raw   = generate(model, word2idx, idx2word,
                             seed_text=seed,
                             length=min(response_length, max_length),
                             temperature=temperature)
            reply = format_response(raw)
            print(f"5x2T: {reply}\n")


# ---------------- ENTRY POINT ----------------
if __name__ == "__main__":
    import sys
    if len(sys.argv) > 1 and sys.argv[1] == "chat":
        chat()
    else:
        model, word2idx, idx2word = train()
        print("\nStarting chat...\n")
        chat(model, word2idx, idx2word)