"""
neural_network.py — Real PyTorch neural network built from scratch.

KEY CHANGES:
  - Every item ingested is IMMEDIATELY written to knowledge.jsonl
  - On startup, knowledge.jsonl is read back → data_buffer is restored
  - Model checkpoint auto-saves every 30 training epochs
  - training_stats.json written after every training step (human-readable)
"""

import torch
import torch.nn as nn
import torch.optim as optim
import threading
_TEXT_LOCK = threading.Lock()
import numpy as np
import os
import json
import re
from datetime import datetime, UTC
from collections import Counter, defaultdict

# ─── FILES ON DISK ────────────────────────────────────────────────────────────
KNOWLEDGE_FILE   = 'knowledge.jsonl'      # Every article/text the AI has seen
CHECKPOINT_FILE  = 'model_checkpoint.pt'  # PyTorch weights + optimizer state
STATS_FILE       = 'training_stats.json'  # Human-readable live stats

# ─── CATEGORIES ───────────────────────────────────────────────────────────────
CATEGORIES = ['technology', 'science', 'world', 'sports',
              'business', 'health', 'entertainment', 'other']

# ─── VOCABULARY ───────────────────────────────────────────────────────────────
STOPWORDS = {
    'a','an','the','is','it','in','on','at','to','for','of','and','or','but',
    'was','are','were','be','been','have','has','had','do','does','did','will',
    'would','could','should','may','might','that','this','these','those','with',
    'from','by','as','not','also','than','then','so','if','when','what','how',
    'who','which','its','their','our','your','my','his','her','we','they','he',
    'she','you','i','me','him','us','them','said','says','new','one','two',
}

class Vocabulary:
    def __init__(self, max_size=10000):
        self.word2idx   = {'<PAD>': 0, '<UNK>': 1}
        self.idx2word   = {0: '<PAD>', 1: '<UNK>'}
        self.word_counts = Counter()
        self.max_size   = max_size
        self.is_built   = False

    def update(self, text: str):
        self.word_counts.update(self._tokenize(text))

    def build(self):
        top = self.word_counts.most_common(self.max_size - 2)
        self.word2idx = {'<PAD>': 0, '<UNK>': 1}
        self.idx2word = {0: '<PAD>', 1: '<UNK>'}
        for i, (word, _) in enumerate(top):
            idx = i + 2
            self.word2idx[word] = idx
            self.idx2word[idx]  = word
        self.is_built = True

    def encode(self, text: str, max_len: int = 64) -> list:
        words = self._tokenize(text)[:max_len]
        ids   = [self.word2idx.get(w, 1) for w in words]
        ids  += [0] * (max_len - len(ids))
        return ids

    def _tokenize(self, text: str) -> list:
        text = text.lower()
        text = re.sub(r'[^\w\s]', ' ', text)
        return [w for w in text.split() if w not in STOPWORDS and len(w) > 2]

    def __len__(self):
        return len(self.word2idx)


# ─── MODEL ────────────────────────────────────────────────────────────────────
class TextClassifier(nn.Module):
    def __init__(self, vocab_size=10002, embed_dim=64,
                 hidden=[256, 128, 64], num_classes=8):
        super().__init__()
        self.embed_dim   = embed_dim
        self.hidden_dims = hidden
        self.num_classes = num_classes

        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        nn.init.normal_(self.embedding.weight, 0, 0.1)

        layers = []
        in_dim = embed_dim
        for h in hidden:
            layers += [nn.Linear(in_dim, h), nn.LayerNorm(h),
                       nn.ReLU(), nn.Dropout(0.25)]
            in_dim = h
        layers.append(nn.Linear(in_dim, num_classes))
        self.net = nn.Sequential(*layers)

        self._activations = {}
        self._register_hooks()

    def _register_hooks(self):
        def make_hook(name):
            def hook(module, inp, out):
                if isinstance(out, torch.Tensor):
                    v = out.detach().float()
                    if v.dim() > 1:
                        v = v.mean(0)
                    self._activations[name] = v[:32].tolist()
            return hook
        for i, layer in enumerate(self.net):
            layer.register_forward_hook(make_hook(f'net.{i}'))

    def forward(self, x):
        emb    = self.embedding(x)
        mask   = (x != 0).float().unsqueeze(-1)
        pooled = (emb * mask).sum(1) / mask.sum(1).clamp(min=1)
        return self.net(pooled)

    def get_activations(self) -> dict:
        return dict(self._activations)

    def get_weight_info(self) -> dict:
        info = {}
        for name, param in self.named_parameters():
            if 'weight' in name and param.dim() == 2:
                w    = param.detach().float().numpy()
                r, c = min(w.shape[0], 16), min(w.shape[1], 16)
                info[name] = {
                    'shape':    list(w.shape),
                    'mean_abs': float(np.mean(np.abs(w))),
                    'std':      float(np.std(w)),
                    'sample':   w[:r, :c].tolist(),
                }
        return info


# ─── LIVING NETWORK ───────────────────────────────────────────────────────────
class LivingNetwork:
    """
    The brain. Wraps the PyTorch model with:
      - knowledge.jsonl  → persistent record of everything it has read
      - model_checkpoint.pt → saved weights (restored on restart)
      - training_stats.json → live stats readable by the UI
    """

    AUTO_SAVE_EVERY = 30   # Save checkpoint every N training epochs

    def __init__(self):
        self.vocab      = Vocabulary()
        self.model      = TextClassifier()
        self.optimizer  = optim.Adam(self.model.parameters(), lr=0.001, weight_decay=1e-5)
        self.scheduler  = optim.lr_scheduler.ReduceLROnPlateau(
            self.optimizer, mode='min', patience=20, factor=0.5, min_lr=1e-5)
        self.criterion  = nn.CrossEntropyLoss()

        self.epoch           = 0
        self.total_samples   = 0
        self.data_buffer     = []   # (text, label_int) — in-memory training pool
        self.loss_history    = []
        self.acc_history     = []
        self.category_counts = defaultdict(int)
        self.knowledge_count = 0    # Total articles ever ingested

        self.stats = {
            'epoch':         0,
            'loss':          '—',
            'accuracy':      '—',
            'total_samples': 0,
            'lr':            0.001,
            'buffer_size':   0,
            'knowledge_count': 0,
            'last_text':     '(nothing yet)',
            'vocab_size':    2,
            'status':        'idle',
        }

        # Load checkpoint first, then restore knowledge buffer
        self._load_checkpoint()
        self._load_knowledge()

    # ── KNOWLEDGE FILE ────────────────────────────────────────────────────────
    def _write_knowledge(self, text: str, category: str, source: str = 'unknown'):
        """Append one learned item to knowledge.jsonl immediately."""
        record = {
            'text':      text,
            'category':  category,
            'source':    source,
            'timestamp': datetime.now(UTC).isoformat(),
            'epoch_at_ingestion': self.epoch,
        }
        try:
            with open(KNOWLEDGE_FILE, 'a', encoding='utf-8') as f:
                f.write(json.dumps(record, ensure_ascii=False) + '\n')
            self.knowledge_count += 1
        except Exception:
            pass

    def _load_knowledge(self):
        """On startup: read knowledge.jsonl and rebuild data_buffer + vocab."""
        if not os.path.exists(KNOWLEDGE_FILE):
            return
        loaded = 0
        try:
            with open(KNOWLEDGE_FILE, 'r', encoding='utf-8') as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        rec   = json.loads(line)
                        text  = rec.get('text', '')
                        cat   = rec.get('category', 'other')
                        label = CATEGORIES.index(cat) if cat in CATEGORIES else 7
                        if len(text) > 20:
                            self.data_buffer.append((text, label))
                            self.vocab.update(text)
                            self.category_counts[cat] += 1
                            loaded += 1
                    except Exception:
                        continue
            self.knowledge_count = loaded
            if loaded > 0:
                self.vocab.build()
            # Trim buffer if huge
            if len(self.data_buffer) > 5000:
                self.data_buffer = self.data_buffer[-4000:]
        except Exception:
            pass
        self.stats['knowledge_count'] = self.knowledge_count
        self.stats['buffer_size']     = len(self.data_buffer)
        self.stats['vocab_size']      = len(self.vocab)

    def get_knowledge_file_stats(self) -> dict:
        """Return stats about the knowledge file for the UI."""
        if not os.path.exists(KNOWLEDGE_FILE):
            return {'exists': False, 'lines': 0, 'size_kb': 0}
        size = os.path.getsize(KNOWLEDGE_FILE)
        lines = 0
        try:
            with open(KNOWLEDGE_FILE, 'r', encoding='utf-8') as f:
                lines = sum(1 for l in f if l.strip())
        except Exception:
            pass
        return {'exists': True, 'lines': lines, 'size_kb': round(size / 1024, 1)}

    def get_recent_knowledge(self, n: int = 20) -> list:
        """Return last N items from knowledge.jsonl for display."""
        if not os.path.exists(KNOWLEDGE_FILE):
            return []
        lines = []
        try:
            with open(KNOWLEDGE_FILE, 'r', encoding='utf-8') as f:
                all_lines = [l.strip() for l in f if l.strip()]
            for line in reversed(all_lines[-n:]):
                try:
                    lines.append(json.loads(line))
                except Exception:
                    pass
        except Exception:
            pass
        return lines

    # ── INGEST ────────────────────────────────────────────────────────────────
    def ingest(self, text: str, category: str, source: str = 'unknown'):
        """
        Add text to training buffer AND write to knowledge.jsonl immediately.
        This is how the AI 'remembers' what it has learned.
        """
        cleaned = text.strip()
        if len(cleaned) < 20:
            return
        label = CATEGORIES.index(category) if category in CATEGORIES else 7

        # ① Write to disk first — never lose this
        self._write_knowledge(cleaned, category, source)

        # ② Add to in-memory training buffer
        self.data_buffer.append((cleaned, label))
        self.vocab.update(cleaned)
        self.category_counts[category] += 1

        # Rebuild vocab every 25 items
        if len(self.data_buffer) % 25 == 0:
            self.vocab.build()

        # Keep buffer bounded (disk has the full history)
        if len(self.data_buffer) > 5000:
            self.data_buffer = self.data_buffer[-4000:]

        self.stats.update({
            'buffer_size':     len(self.data_buffer),
            'vocab_size':      len(self.vocab),
            'knowledge_count': self.knowledge_count,
        })

    # ── TRAINING ──────────────────────────────────────────────────────────────
    def train_step(self, batch_size: int = 32) -> float | None:
        if len(self.data_buffer) < batch_size or not self.vocab.is_built:
            return None
        with _TEXT_LOCK:
            self.model.train()
            idx   = np.random.choice(len(self.data_buffer), batch_size, replace=False)
            batch = [self.data_buffer[i] for i in idx]
            texts, labels = zip(*batch)

            x = torch.tensor([self.vocab.encode(t) for t in texts], dtype=torch.long)
            y = torch.tensor(list(labels), dtype=torch.long)

            self.model.zero_grad(set_to_none=True)
            logits = self.model(x)
            loss   = self.criterion(logits, y)
            loss.backward()
            for p in self.model.parameters():
                if p.grad is not None:
                    p.grad.data.clamp_(-1.0, 1.0)
            self.optimizer.step()
            loss_val = loss.detach().item()
            acc      = (logits.detach().argmax(1) == y).float().mean().item()

        self.epoch        += 1
        self.total_samples += batch_size
        self.scheduler.step(loss_val)

        self.loss_history.append(round(loss_val, 5))
        self.acc_history.append(round(acc, 4))
        if len(self.loss_history) > 500:
            self.loss_history = self.loss_history[-500:]
            self.acc_history  = self.acc_history[-500:]

        self.stats.update({
            'epoch':           self.epoch,
            'loss':            round(loss_val, 4),
            'accuracy':        round(acc * 100, 1),
            'total_samples':   self.total_samples,
            'lr':              round(self.optimizer.param_groups[0]['lr'], 7),
            'buffer_size':     len(self.data_buffer),
            'last_text':       texts[0][:120],
            'vocab_size':      len(self.vocab),
            'knowledge_count': self.knowledge_count,
        })

        # Auto-save checkpoint every N epochs
        if self.epoch % self.AUTO_SAVE_EVERY == 0:
            self.save_checkpoint()

        # Always write stats file so UI can read without waiting
        self._write_stats_file()

        return loss_val

    def train_n_steps(self, n: int = 50) -> dict:
        losses = []
        for _ in range(n):
            l = self.train_step()
            if l is not None:
                losses.append(l)
        return {
            'steps':    len(losses),
            'avg_loss': round(sum(losses) / len(losses), 5) if losses else None,
        }

    # ── INFERENCE ─────────────────────────────────────────────────────────────
    def predict(self, text: str) -> dict:
        if not self.vocab.is_built or not text.strip():
            return {'error': 'Model not ready — start the network and let it train first'}
        self.model.eval()
        with torch.no_grad():
            x      = torch.tensor([self.vocab.encode(text)], dtype=torch.long)
            logits = self.model(x)
            probs  = torch.softmax(logits, dim=1)[0].tolist()
            pred   = int(logits.argmax(1).item())
        return {
            'prediction': CATEGORIES[pred],
            'confidence': round(probs[pred] * 100, 1),
            'all_probs':  {c: round(p * 100, 2) for c, p in zip(CATEGORIES, probs)},
        }

    # ── VIZ STATE ─────────────────────────────────────────────────────────────
    def get_viz_state(self) -> dict:
        self.model.eval()
        with torch.no_grad():
            dummy = torch.zeros(1, 64, dtype=torch.long)
            self.model(dummy)
        return {
            'layer_sizes':    [self.model.embed_dim] + self.model.hidden_dims + [self.model.num_classes],
            'activations':    self.model.get_activations(),
            'weights':        self.model.get_weight_info(),
            'loss_history':   self.loss_history[-100:],
            'acc_history':    self.acc_history[-100:],
            'stats':          self.stats,
            'category_counts':dict(self.category_counts),
        }

    # ── PERSISTENCE ───────────────────────────────────────────────────────────
    def save_checkpoint(self):
        try:
            torch.save({
                'model':           self.model.state_dict(),
                'optimizer':       self.optimizer.state_dict(),
                'epoch':           self.epoch,
                'total_samples':   self.total_samples,
                'loss_history':    self.loss_history,
                'acc_history':     self.acc_history,
                'vocab_word2idx':  self.vocab.word2idx,
                'category_counts': dict(self.category_counts),
                'stats':           self.stats,
            }, CHECKPOINT_FILE)
            return True
        except Exception:
            return False

    def _load_checkpoint(self):
        if not os.path.exists(CHECKPOINT_FILE):
            return False
        try:
            ck = torch.load(CHECKPOINT_FILE, map_location='cpu')
            self.model.load_state_dict(ck['model'])
            self.optimizer.load_state_dict(ck['optimizer'])
            self.epoch           = ck.get('epoch', 0)
            self.total_samples   = ck.get('total_samples', 0)
            self.loss_history    = ck.get('loss_history', [])
            self.acc_history     = ck.get('acc_history', [])
            self.category_counts = defaultdict(int, ck.get('category_counts', {}))
            self.stats           = ck.get('stats', self.stats)
            w2i = ck.get('vocab_word2idx', {})
            if w2i:
                self.vocab.word2idx  = w2i
                self.vocab.idx2word  = {v: k for k, v in w2i.items()}
                self.vocab.is_built  = len(w2i) > 2
            return True
        except Exception:
            return False

    def _write_stats_file(self):
        """Write human-readable stats to training_stats.json for easy debugging."""
        try:
            with open(STATS_FILE, 'w') as f:
                json.dump({**self.stats, 'loss_last10': self.loss_history[-10:]}, f, indent=2)
        except Exception:
            pass