use candle_core::{Result, Tensor, IndexOp};
use candle_nn::{Module, VarBuilder};
use super::attention::Attention;
use super::moe::MoeBlock;
use super::mlp::Mlp;
use super::config::TransformerConfig;
use super::ternary_linear::TernaryLinear;
use super::anastomosis::{AnastomosisLayer, capture_cosim};

pub struct Block {
    attention: Attention,
    moe: Option<MoeBlock>,
    mlp: Option<Mlp>,
    ln1: candle_nn::LayerNorm,
    ln2: candle_nn::LayerNorm,
}

impl Block {
    pub fn new(config: &TransformerConfig, vb: VarBuilder, threshold: f32) -> Result<Self> {
        let attention = Attention::new(config.hidden_size, config.num_heads, vb.pp("attn"), threshold)?;

        let (moe, mlp) = if config.num_experts > 0 {
            let moe = MoeBlock::new(config.hidden_size, config.num_experts, vb.pp("moe"), threshold)?;
            (Some(moe), None)
        } else {
            let mlp = Mlp::new(config.hidden_size, config.hidden_size * 4, vb.pp("mlp"), threshold)?;
            (None, Some(mlp))
        };

        let ln1 = candle_nn::layer_norm(config.hidden_size, 1e-5, vb.pp("ln1"))?;
        let ln2 = candle_nn::layer_norm(config.hidden_size, 1e-5, vb.pp("ln2"))?;
        Ok(Self { attention, moe, mlp, ln1, ln2 })
    }

    /// Dual-stream constructor: stream-specific parts (attn, gate, ln) from `vb_stream`,
    /// shared expert FFN weights from `vb_experts`. Both streams referencing the same
    /// `vb_experts` prefix ensures shared tensor variables in the VarMap.
    pub fn new_stream(config: &TransformerConfig, vb_stream: VarBuilder, vb_experts: VarBuilder, threshold: f32) -> Result<Self> {
        let attention = Attention::new(config.hidden_size, config.num_heads, vb_stream.pp("attn"), threshold)?;

        let (moe, mlp) = if config.num_experts > 0 {
            let moe = MoeBlock::new_stream(
                config.hidden_size, config.num_experts,
                vb_stream.pp("moe"), vb_experts, threshold,
            )?;
            (Some(moe), None)
        } else {
            let mlp = Mlp::new(config.hidden_size, config.hidden_size * 4, vb_stream.pp("mlp"), threshold)?;
            (None, Some(mlp))
        };

        let ln1 = candle_nn::layer_norm(config.hidden_size, 1e-5, vb_stream.pp("ln1"))?;
        let ln2 = candle_nn::layer_norm(config.hidden_size, 1e-5, vb_stream.pp("ln2"))?;
        Ok(Self { attention, moe, mlp, ln1, ln2 })
    }

    pub fn prepare_inference(&self) -> Result<()> {
        self.attention.prepare_inference()?;
        if let Some(moe) = &self.moe { moe.prepare_inference()?; }
        if let Some(mlp) = &self.mlp { mlp.prepare_inference()?; }
        Ok(())
    }

    pub fn freeze_ttl(&self, steps: usize) {
        if let Some(moe) = &self.moe { moe.freeze_ttl(steps); }
    }

    pub fn clear_kv_cache(&self) {
        self.attention.clear_kv_cache();
    }

    // Training / full-sequence forward.
    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
        let x = (x + self.attention.forward(&self.ln1.forward(x)?)?)?;
        self.feed_forward(x)
    }

    // Prefill: process full prompt, populate KV-cache.
    pub fn forward_and_cache(&self, x: &Tensor) -> Result<Tensor> {
        let x = (x + self.attention.forward_and_cache(&self.ln1.forward(x)?)?)?;
        self.feed_forward(x)
    }

    // Decode: single new token, reads from KV-cache.
    pub fn forward_decode(&self, x: &Tensor) -> Result<Tensor> {
        let x = (x + self.attention.forward_decode(&self.ln1.forward(x)?)?)?;
        self.feed_forward(x)
    }

    fn feed_forward(&self, x: Tensor) -> Result<Tensor> {
        if let Some(moe) = &self.moe {
            Ok((&x + moe.forward(&self.ln2.forward(&x)?)?)?)
        } else if let Some(mlp) = &self.mlp {
            Ok((&x + mlp.forward(&self.ln2.forward(&x)?)?)?)
        } else {
            Ok(x)
        }
    }
}

pub struct Transformer {
    embedding: candle_nn::Embedding,
    blocks: Vec<Block>,
    /// Stream B blocks — present only in dual-stream (cord) mode.
    blocks_b: Option<Vec<Block>>,
    /// Anastomosis gates at Fibonacci-indexed layers. Empty in single-stream mode.
    anastomosis: Vec<(usize, AnastomosisLayer)>,
    ln_f: candle_nn::LayerNorm,
    lm_head: TernaryLinear,
    config: TransformerConfig,
}

impl Transformer {
    pub fn new(config: &TransformerConfig, vb: VarBuilder) -> Result<Self> {
        let embedding = candle_nn::embedding(config.vocab_size, config.hidden_size, vb.pp("embed"))?;

        let mut blocks = Vec::new();
        let mut blocks_b: Option<Vec<Block>> = None;
        let mut anastomosis: Vec<(usize, AnastomosisLayer)> = Vec::new();

        let vb_blocks = vb.pp("blocks");

        if config.num_streams >= 2 {
            let mut blocks_b_vec = Vec::new();
            for i in 0..config.num_layers {
                let thr = config.layer_threshold(i);
                let vb_layer = vb_blocks.pp(i);
                let vb_experts = vb_layer.pp("experts");
                blocks.push(Block::new_stream(config, vb_layer.pp("stream_a"), vb_experts.clone(), thr)?);
                blocks_b_vec.push(Block::new_stream(config, vb_layer.pp("stream_b"), vb_experts, thr)?);
            }
            blocks_b = Some(blocks_b_vec);

            for (k, &fusion_idx) in config.fusion_layers.iter().enumerate() {
                anastomosis.push((fusion_idx, AnastomosisLayer::new(config.hidden_size, vb.pp("anastomosis").pp(k))?));
            }
        } else {
            for i in 0..config.num_layers {
                let thr = config.layer_threshold(i);
                blocks.push(Block::new(config, vb_blocks.pp(i), thr)?);
            }
        }

        let ln_f = candle_nn::layer_norm(config.hidden_size, 1e-5, vb.pp("ln_f"))?;
        let threshold = config.threshold;
        let lm_head = TernaryLinear::new(config.hidden_size, config.vocab_size, false, threshold, vb.pp("lm_head"))?;

        Ok(Self { embedding, blocks, blocks_b, anastomosis, ln_f, lm_head, config: config.clone() })
    }

    /// Pre-ternarize all layer weights + pre-warm the lm_head.
    /// Call once after loading a checkpoint for inference; eliminates per-step re-quantization.
    pub fn prepare_inference(&self) -> Result<()> {
        for block in &self.blocks { block.prepare_inference()?; }
        if let Some(bb) = &self.blocks_b {
            for block in bb { block.prepare_inference()?; }
        }
        self.lm_head.prepare_inference()
    }

    /// Freeze TTL logit modifiers on a specific transformer layer for `steps` update() calls.
    /// Called by train_bible on gradient-norm burst detection to allow gate learning.
    pub fn freeze_ttl_layer(&self, layer_idx: usize, steps: usize) {
        if let Some(block) = self.blocks.get(layer_idx) {
            block.freeze_ttl(steps);
        }
        if let Some(bb) = &self.blocks_b {
            if let Some(block) = bb.get(layer_idx) {
                block.freeze_ttl(steps);
            }
        }
    }

    /// Reset all per-layer KV-caches. Call before each new generation sequence.
    pub fn clear_kv_cache(&self) {
        for block in &self.blocks { block.clear_kv_cache(); }
        if let Some(bb) = &self.blocks_b {
            for block in bb { block.clear_kv_cache(); }
        }
    }

    // Training / evaluation forward — no KV-cache side effects.
    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
        let embed = self.embedding.forward(x)?;

        if let Some(blocks_b) = &self.blocks_b {
            let mut h_a = embed.clone();
            let mut h_b = embed;

            for i in 0..self.blocks.len() {
                h_a = self.blocks[i].forward(&h_a)?;
                h_b = blocks_b[i].forward(&h_b)?;

                capture_cosim(&h_a, &h_b);

                if let Some((_, anast)) = self.anastomosis.iter().find(|(idx, _)| *idx == i) {
                    let (new_a, new_b) = anast.forward(&h_a, &h_b)?;
                    h_a = new_a;
                    h_b = new_b;
                }
            }

            let h_merged = (&h_a + &h_b)?.affine(0.5, 0.0)?;
            let h_merged = self.ln_f.forward(&h_merged)?;
            self.lm_head.forward(&h_merged)
        } else {
            let mut x = embed;
            for block in &self.blocks { x = block.forward(&x)?; }
            x = self.ln_f.forward(&x)?;
            self.lm_head.forward(&x)
        }
    }

    /// Prefill: process full prompt through all layers, populate KV-cache in each attention block.
    /// Returns full logit tensor [1, seq_len, vocab]; caller takes the last position.
    pub fn forward_prefill(&self, x: &Tensor) -> Result<Tensor> {
        let embed = self.embedding.forward(x)?;

        if let Some(blocks_b) = &self.blocks_b {
            let mut h_a = embed.clone();
            let mut h_b = embed;

            for i in 0..self.blocks.len() {
                h_a = self.blocks[i].forward_and_cache(&h_a)?;
                h_b = blocks_b[i].forward_and_cache(&h_b)?;

                if let Some((_, anast)) = self.anastomosis.iter().find(|(idx, _)| *idx == i) {
                    let (new_a, new_b) = anast.forward(&h_a, &h_b)?;
                    h_a = new_a;
                    h_b = new_b;
                }
            }

            let h_merged = (&h_a + &h_b)?.affine(0.5, 0.0)?;
            let h_merged = self.ln_f.forward(&h_merged)?;
            self.lm_head.forward(&h_merged)
        } else {
            let mut h = embed;
            for block in &self.blocks { h = block.forward_and_cache(&h)?; }
            h = self.ln_f.forward(&h)?;
            self.lm_head.forward(&h)
        }
    }

    /// Decode: single new token. Uses KV-cache — O(1) per layer instead of O(seq_len).
    /// Position is tracked internally via the KV-cache length in each attention block.
    /// Returns logits [1, 1, vocab]; caller takes [0, 0].
    pub fn forward_decode(&self, token: u32, dev: &candle_core::Device) -> Result<Tensor> {
        let x = Tensor::new(&[token], dev)?
            .unsqueeze(0)?.to_dtype(candle_core::DType::U32)?;
        let embed = self.embedding.forward(&x)?;

        if let Some(blocks_b) = &self.blocks_b {
            let mut h_a = embed.clone();
            let mut h_b = embed;

            for i in 0..self.blocks.len() {
                h_a = self.blocks[i].forward_decode(&h_a)?;
                h_b = blocks_b[i].forward_decode(&h_b)?;

                if let Some((_, anast)) = self.anastomosis.iter().find(|(idx, _)| *idx == i) {
                    let (new_a, new_b) = anast.forward(&h_a, &h_b)?;
                    h_a = new_a;
                    h_b = new_b;
                }
            }

            let h_merged = (&h_a + &h_b)?.affine(0.5, 0.0)?;
            let h_merged = self.ln_f.forward(&h_merged)?;
            self.lm_head.forward(&h_merged)
        } else {
            let mut h = embed;
            for block in &self.blocks { h = block.forward_decode(&h)?; }
            h = self.ln_f.forward(&h)?;
            self.lm_head.forward(&h)
        }
    }

    pub fn generate(&self, tokenizer: &super::super::tokenizer::BpeTokenizer, prompt: &str, max_new_tokens: usize) -> String {
        let mut tokens = tokenizer.encode(prompt);
        if tokens.is_empty() { return "".to_string(); }

        let dev = candle_core::Device::Cpu;
        let mut rng = rand::thread_rng();

        for _ in 0..max_new_tokens {
            let start_idx = tokens.len().saturating_sub(self.config.max_seq_len);
            let context = &tokens[start_idx..];

            let input = Tensor::new(&context[..], &dev).unwrap()
                .unsqueeze(0).unwrap().to_dtype(candle_core::DType::U32).unwrap();
            let logits = self.forward(&input).unwrap();

            let dims = logits.dims();
            let last_logits = logits.i((0, dims[1] - 1)).unwrap();

            let probs = candle_nn::ops::softmax(&last_logits, 0).unwrap();
            let pr = probs.to_vec1::<f32>().unwrap();

            use rand::distributions::{Distribution, WeightedIndex};
            let dist = WeightedIndex::new(&pr).unwrap();
            let next_token_idx = dist.sample(&mut rng) as u32;
            tokens.push(next_token_idx);
        }

        tokenizer.decode(&tokens)
    }
}
