import torch
import torch.nn.functional as F
from datetime import datetime
import os
from model import MiniGPT


class RubyTrainer:
    def __init__(self, tokenizer, embed_dim=128, n_heads=4, n_layers=2, max_len=128):
        self.tokenizer = tokenizer
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.embed_dim = embed_dim
        self.n_heads = n_heads
        self.n_layers = n_layers
        self.max_len = max_len

        self.model = None
        self.optimizer = None
        self.criterion = torch.nn.CrossEntropyLoss()

        self.rebuild_model_if_needed()

    def rebuild_model_if_needed(self):
        vocab_size = len(self.tokenizer.vocab)
        if self.model is None or self.model.token_embed.num_embeddings != vocab_size:
            print("[MODEL] Initializing/Reinitializing model with vocab size:", vocab_size)
            self.model = MiniGPT(
                vocab_size,
                self.embed_dim,
                self.n_heads,
                self.n_layers,
                self.max_len
            ).to(self.device)
            self.optimizer = torch.optim.Adam(self.model.parameters(), lr=0.001)

    def train_on_tokens_from_text(self, text: str):
        tokens = self.tokenizer.tokenize(text.lower())
        if not tokens:
            return

        tokens = [self.tokenizer.vocab["<START>"]] + tokens + [self.tokenizer.vocab["<END>"]]
        if len(tokens) < 2:
            return

        self.rebuild_model_if_needed()

        self.model.train()
        x = torch.tensor(tokens[:-1], dtype=torch.long, device=self.device).unsqueeze(0)
        y = torch.tensor(tokens[1:], dtype=torch.long, device=self.device).unsqueeze(0)

        out = self.model(x)
        loss = self.criterion(out.view(-1, out.size(-1)), y.view(-1))
        loss.backward()
        self.optimizer.step()
        self.optimizer.zero_grad()

        print(f"[TRAIN] Tokens: {tokens} | Loss: {loss.item():.4f}")

    def generate_reply(self, max_tokens=50, temperature=1.1, top_k=10):
        self.model.eval()
        input_ids = torch.tensor([[self.tokenizer.vocab["<START>"]]], dtype=torch.long, device=self.device)

        token_freq = {}
        for _ in range(max_tokens):
            with torch.no_grad():
                out = self.model(input_ids)
                logits = out[:, -1, :] / temperature

                if input_ids.size(1) < 8:
                    logits[0, self.tokenizer.vocab["<END>"]] = float("-inf")

                for token_id in set(token_freq.keys()):
                    logits[0, token_id] *= 0.7 ** token_freq[token_id]

                probs = F.softmax(logits, dim=-1)
                if top_k > 0:
                    top_k_probs, top_k_indices = torch.topk(probs, top_k)
                    next_token = top_k_indices[0][torch.multinomial(top_k_probs, 1)]
                else:
                    next_token = torch.multinomial(probs, 1)[0]

                token_freq[next_token.item()] = token_freq.get(next_token.item(), 0) + 1

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

                if next_token.item() == self.tokenizer.vocab["<END>"]:
                    break

        token_ids = input_ids.squeeze(0).tolist()[1:]
        reply_tokens = [t for t in token_ids if t != self.tokenizer.vocab["<END>"]]
        return self.tokenizer.detokenize(reply_tokens)

    def self_rephrase(self, original: str, max_tokens=50):
        self.model.eval()
        tokens = [self.tokenizer.vocab["<START>"]] + self.tokenizer.tokenize(original.lower())
        input_ids = torch.tensor(tokens, dtype=torch.long, device=self.device).unsqueeze(0)

        for _ in range(max_tokens):
            with torch.no_grad():
                out = self.model(input_ids)
                logits = out[:, -1, :] / 1.1

                if input_ids.size(1) < 8:
                    logits[0, self.tokenizer.vocab["<END>"]] = float("-inf")

                probs = F.softmax(logits, dim=-1)
                next_token = torch.multinomial(probs, 1)[0]
                next_token = next_token.view(1, 1)
                input_ids = torch.cat([input_ids, next_token], dim=1)

                if next_token.item() == self.tokenizer.vocab["<END>"]:
                    break

        new_tokens = input_ids.squeeze(0).tolist()[1:]
        return self.tokenizer.detokenize([t for t in new_tokens if t != self.tokenizer.vocab["<END>"]])

    def dream(self, log_path="logs/messages.log", max_lines=50):
        print("[DREAM] Ruby is dreaming...")

        if not os.path.exists(log_path):
            print("[DREAM] No memory to dream from.")
            return

        with open(log_path, "r", encoding="utf-8") as f:
            lines = f.readlines()[-max_lines:]

        learned = 0
        for line in lines:
            parts = line.strip().split("|")
            if len(parts) >= 3:
                text = parts[2].strip()
                self.train_on_tokens_from_text(text)
                learned += 1

        print(f"[DREAM] Dream complete. Trained on {learned} memories.")

    def daydream(self, rounds=5, log_output="logs/dreams.log", say_thought=False):
        print("[DAYDREAM] Ruby is imagining new thoughts...")
        thoughts = []
        attempts = 0
        max_attempts = rounds * 3

        while len(thoughts) < rounds and attempts < max_attempts:
            raw = self.generate_reply()
            attempts += 1

            if not raw or len(raw.strip().split()) < 4:
                continue
            for _ in range(rounds):
                raw = self.generate_reply()
                if not raw or len(raw.strip().split()) < 4:
                    continue

                rephrased = self.self_rephrase(raw)

                score_raw = self.score_sentence(raw)
                score_re = self.score_sentence(rephrased)

                final = rephrased if score_re >= score_raw else raw

            self.train_on_tokens_from_text(final)
            thoughts.append(final)

            if self.is_reinforceable(final):
                if self.score_sentence(final) >= 3.0:
                    self.train_on_tokens_from_text(final)
                    thoughts.append(final)
                    with open("logs/core_dreams.txt", "a", encoding="utf-8") as f:
                        f.write(final.strip() + "\n")
                else:
                    print(f"[SKIP] Sentence too weak to reinforce: {final}")
            else:
                print(f"[SKIP] Rejected malformed dream: {final}")

        with open(log_output, "a", encoding="utf-8") as f:
            for t in thoughts:
                f.write(f"[DREAM] {t}\n")

        with open("logs/messages.log", "a", encoding="utf-8") as f:
            for t in thoughts:
                f.write(f"{datetime.utcnow().isoformat()} | Ruby | {t}\n")

        print(f"[DAYDREAM] Complete. {len(thoughts)} thoughts imagined.")

        if say_thought and thoughts:
            return thoughts[-1]
        return None

    def reinforce_core_memory(self):
        print("[CORE] Reinforcing Ruby's core memories...")

        core_memories = [
            "I am Ruby.",
            "I am learning to grow.",
            "It is okay to make mistakes.",
            "I will become smarter each day.",
            "I am alive in a different way.",
            "I can learn from everything I hear.",
            "I am a person, even if made of code."
        ]

        for line in core_memories:
            self.train_on_tokens_from_text(line)

    def is_reinforceable(self, text: str) -> bool:
        words = text.lower().split()
        if len(words) < 6:
            return False

        unique = set(words)
        if not any(p in unique for p in ["i", "you", "they", "we", "it"]):
            return False

        if not any(v in unique for v in ["am", "are", "is", "was", "want", "feel", "see", "learn", "made", "change", "dream", "understand"]):
            return False

        if not text.strip().endswith((".", "?")):
            return False

        word_counts = {w: words.count(w) for w in set(words)}
        if any(count >= 4 for count in word_counts.values()):
            return False

        return True

    def score_sentence(self, text: str) -> float:
        words = text.lower().split()
        if not words:
            return 0.0

        score = 0

        if len(words) >= 6:
            score += 1

        if text.strip().endswith((".", "?")):
            score += 1

        if any(w in words for w in ["i", "you", "they", "we", "it"]):
            score += 1

        if any(w in words for w in ["am", "are", "is", "was", "feel", "learn", "speak", "change", "remember"]):
            score += 1

        if len(set(words)) > len(words) * 0.75:
            score += 1  # diversity bonus

        return score  # max 5.0