//! CPU reference implementation of Falcon-H1 inference mathematics, not a CUDA backend.
//! Checkpoint tensors are loaded from the simple FH1W/f32 pack produced by tools/export.py.
//! Intended as a readable, testable baseline; not optimized for 0.5B CPU throughput.
use std::collections::HashMap;
use std::fs::File;
use std::io::{self, Read};

#[derive(Clone, Debug)]
pub struct Config {
    pub vocab: usize, pub hidden: usize, pub layers: usize, pub heads: usize,
    pub kv_heads: usize, pub head_dim: usize, pub mamba_heads: usize,
    pub mamba_dim: usize, pub state: usize, pub groups: usize, pub conv: usize,
    pub ff: usize, pub rope_theta: f32, pub eps: f32,
    pub embedding_multiplier: f32, pub key_multiplier: f32,
    pub attention_in_multiplier: f32, pub attention_out_multiplier: f32,
    pub ssm_in_multiplier: f32, pub ssm_out_multiplier: f32,
    pub ssm_multipliers: [f32;5], pub mlp_multipliers: [f32;2],
    pub lm_head_multiplier: f32,
}
impl Config {
    pub fn falcon_h1_05b() -> Self { Self {
        vocab:32784, hidden:1024, layers:36, heads:8, kv_heads:2, head_dim:64,
        mamba_heads:24, mamba_dim:1536, state:128, groups:1, conv:4, ff:2048,
        rope_theta:1e11, eps:1e-5,
        embedding_multiplier:5.656854249492381, key_multiplier:0.390625,
        attention_in_multiplier:1., attention_out_multiplier:0.9375,
        ssm_in_multiplier:1.25, ssm_out_multiplier:0.23570226039551587,
        ssm_multipliers:[0.3535533905932738,0.25,0.3535533905932738,0.5,0.3535533905932738],
        mlp_multipliers:[0.8838834764831844,0.5859375], lm_head_multiplier:0.0390625,
    }}
    pub fn demo() -> Self { Self {
        vocab:64, hidden:32, layers:2, heads:4, kv_heads:2, head_dim:8,
        mamba_heads:4, mamba_dim:32, state:8, groups:1, conv:4, ff:64,
        rope_theta:1e11, eps:1e-5,
        embedding_multiplier:1., key_multiplier:1., attention_in_multiplier:1.,
        attention_out_multiplier:1., ssm_in_multiplier:1., ssm_out_multiplier:1.,
        ssm_multipliers:[1.;5], mlp_multipliers:[1.;2], lm_head_multiplier:1.,
    }}
    pub fn validate(&self) -> Result<(), String> {
        if self.heads==0 || self.kv_heads==0 || self.mamba_heads==0 || self.groups==0 || self.conv==0 || self.hidden==0 || self.vocab==0 || self.state==0 { return Err("zero dimension".into()) }
        if self.heads % self.kv_heads != 0 || self.head_dim % 2 != 0 { return Err("invalid GQA/RoPE shape".into()) }
        if self.mamba_dim % self.mamba_heads != 0 || self.mamba_heads % self.groups != 0 { return Err("invalid Mamba head/group shape".into()) }
        // Falcon-H1 attention is narrower than the hidden state: 8*64 != 1024 for 0.5B.
        if self.head_dim==0 || self.mamba_dim==0 || self.ff==0 { return Err("zero projection dimension".into()) }
        Ok(())
    }
}

#[derive(Clone, Debug)]
pub struct Tensor { pub shape: Vec<usize>, pub data: Vec<f32> }
impl Tensor {
    fn len(&self) -> usize { self.data.len() }
}
pub type Weights = HashMap<String,Tensor>;
fn bad(e: &str) -> io::Error { io::Error::new(io::ErrorKind::InvalidData,e) }
fn take_u32(r:&mut impl Read) -> io::Result<u32> { let mut a=[0u8;4]; r.read_exact(&mut a)?; Ok(u32::from_le_bytes(a)) }

/// f32 pack: magic FH1W, u32 count, then each tensor: u16 name_len, name,
/// u8 ndim, u32 dims[ndim], little endian f32 values. No pickle/unsafe code.
pub fn load_pack(path:&str)->io::Result<Weights>{
    let mut f=File::open(path)?; let mut magic=[0u8;4]; f.read_exact(&mut magic)?;
    if &magic!=b"FH1W" { return Err(bad("not a FH1W checkpoint")) }
    let count=take_u32(&mut f)? as usize;
    if count>50000 { return Err(bad("unreasonable tensor count")) }
    let mut weights=HashMap::new();
    for _ in 0..count {
        let mut two=[0u8;2]; f.read_exact(&mut two)?;let n=u16::from_le_bytes(two) as usize;
        if n==0 || n>1024 { return Err(bad("invalid name length")) }
        let mut name=vec![0u8;n];f.read_exact(&mut name)?;
        let name=String::from_utf8(name).map_err(|_|bad("invalid utf8"))?;
        let mut one=[0u8];f.read_exact(&mut one)?;let d=one[0] as usize;
        if d==0 || d>4 { return Err(bad("invalid rank")) }
        let mut shape=vec![];let mut len=1usize;
        for _ in 0..d {let size=take_u32(&mut f)? as usize;
            len=len.checked_mul(size).ok_or_else(||bad("tensor overflow"))?;
            shape.push(size);
        }
        if len>1_000_000_000 {return Err(bad("unreasonable tensor size"))}
        let mut bytes=vec![0u8;len.checked_mul(4).ok_or_else(||bad("size overflow"))?]; f.read_exact(&mut bytes)?;
        let data=bytes.chunks_exact(4).map(|v|f32::from_le_bytes([v[0],v[1],v[2],v[3]])).collect();
        if weights.insert(name,Tensor{shape,data}).is_some(){return Err(bad("duplicate tensor"))}
    }
    Ok(weights)
}

pub fn expected_shapes(c:&Config)->HashMap<String,Vec<usize>>{
    let mut m=HashMap::new();
    m.insert("model.embed_tokens.weight".into(),vec![c.vocab,c.hidden]);
    m.insert("model.final_layernorm.weight".into(),vec![c.hidden]);
    m.insert("lm_head.weight".into(),vec![c.vocab,c.hidden]);
    let cd=c.mamba_dim+2*c.groups*c.state;
    for l in 0..c.layers {
        let p=format!("model.layers.{l}.");
        for n in ["input_layernorm.weight","pre_ff_layernorm.weight"] {m.insert(format!("{p}{n}"),vec![c.hidden]);}
        for (n,r,k) in [
            ("self_attn.q_proj.weight",c.heads*c.head_dim,c.hidden),
            ("self_attn.k_proj.weight",c.kv_heads*c.head_dim,c.hidden),
            ("self_attn.v_proj.weight",c.kv_heads*c.head_dim,c.hidden),
            ("self_attn.o_proj.weight",c.hidden,c.heads*c.head_dim),
            ("feed_forward.gate_proj.weight",c.ff,c.hidden),
            ("feed_forward.up_proj.weight",c.ff,c.hidden),
            ("feed_forward.down_proj.weight",c.hidden,c.ff),
            ("mamba.in_proj.weight",c.mamba_dim+cd+c.mamba_heads,c.hidden),
            ("mamba.out_proj.weight",c.hidden,c.mamba_dim),
        ] {m.insert(format!("{p}{n}"),vec![r,k]);}
        for n in ["mamba.dt_bias","mamba.A_log","mamba.D"]{m.insert(format!("{p}{n}"),vec![c.mamba_heads]);}
        m.insert(format!("{p}mamba.conv1d.weight"),vec![cd,1,c.conv]);
        m.insert(format!("{p}mamba.conv1d.bias"),vec![cd]);
    }
    m
}
pub fn validate_weights(c:&Config,w:&Weights)->Result<(),String>{
    c.validate()?;
    for (k,s) in expected_shapes(c) {
        let t=w.get(&k).ok_or_else(||format!("missing tensor {k}"))?;
        if t.shape!=s { return Err(format!("bad shape for {k}: expected {s:?}, got {:?}",t.shape)) }
        if t.shape.iter().product::<usize>()!=t.len(){return Err(format!("bad data length for {k}"))}
    }
    Ok(())
}
fn silu(x:f32)->f32 { x/(1.+(-x).exp()) }
fn softplus(x:f32)->f32 { if x>20. {x} else {(1.+x.exp()).ln()} }
fn linear(x:&[f32],w:&Tensor)->Vec<f32>{
    let (rows,cols)=(w.shape[0],w.shape[1]);debug_assert_eq!(cols,x.len());
    let mut y=vec![0.;rows];
    for i in 0..rows {let mut v=0f32;for j in 0..cols{v+=w.data[i*cols+j]*x[j];}y[i]=v;}
    y
}
fn norm(x:&[f32],weight:&Tensor,eps:f32)->Vec<f32>{
    let rms=(x.iter().map(|v|v*v).sum::<f32>()/(x.len() as f32)+eps).sqrt();
    x.iter().zip(&weight.data).map(|(a,b)|a/rms*b).collect()
}
fn rope(x:&mut [f32],pos:usize,theta:f32){
    let half=x.len()/2;
    for i in 0..half {
        let angle=(pos as f32)/theta.powf(2.*i as f32/x.len() as f32);
        let (s,c)=angle.sin_cos();let (a,b)=(x[i],x[i+half]);
        x[i]=a*c-b*s; x[i+half]=b*c+a*s;
    }
}
fn add_assign(x:&mut[f32],y:&[f32]) {for (a,b) in x.iter_mut().zip(y){*a+=b}}

/// Reference full-prefix forward. Each layer runs GQA and Mamba-2 from the SAME
/// RMS-normalized activations, then SwiGLU after a second RMSNorm.
/// No KV/SSM caching yet; per-request inference recomputes the prefix.
pub struct Model {pub config:Config, pub weights:Weights}
impl Model {
    pub fn new(config:Config,weights:Weights)->Result<Self,String>{validate_weights(&config,&weights)?;Ok(Self{config,weights})}
    fn t(&self,key:&str)->&Tensor {&self.weights[key]}
    pub fn logits(&self,ids:&[usize])->Result<Vec<f32>,String>{
        let c=&self.config;
        if ids.is_empty(){return Err("input_ids must not be empty".into())}
        if ids.len()>16384{return Err("sequence too long".into())}
        for &i in ids {if i>=c.vocab{return Err(format!("token {i} outside vocabulary"))}}
        let embeddings=self.t("model.embed_tokens.weight");
        let mut x:Vec<Vec<f32>>=ids.iter().map(|id| embeddings.data[id*c.hidden..(id+1)*c.hidden].iter().map(|v|v*c.embedding_multiplier).collect()).collect();
        let n=x.len();let s=c.state;let md=c.mamba_dim;let hd=md/c.mamba_heads;
        let cd=md+2*c.groups*s;
        for layer in 0..c.layers {
            let p=format!("model.layers.{layer}.");
            let normalized:Vec<Vec<f32>>=x.iter().map(|a|norm(a,self.t(&format!("{p}input_layernorm.weight")),c.eps)).collect();
            // Mamba's causal depthwise conv and selective SSM scan.
            let mut conv_history=vec![vec![0f32;c.conv];cd];
            // [head, head_dim, state]
            let mut recurrent=vec![0f32;c.mamba_heads*hd*s];
            let mut mamba_out=vec![vec![0f32;c.hidden];n];
            let in_proj=self.t(&format!("{p}mamba.in_proj.weight"));
            let out_proj=self.t(&format!("{p}mamba.out_proj.weight"));
            let cw=self.t(&format!("{p}mamba.conv1d.weight"));
            let cb=self.t(&format!("{p}mamba.conv1d.bias"));
            let a_log=self.t(&format!("{p}mamba.A_log"));
            let dt_bias=self.t(&format!("{p}mamba.dt_bias"));
            let d_skip=self.t(&format!("{p}mamba.D"));
            let m=c.ssm_multipliers;
            for time in 0..n {
                let input:Vec<f32>=normalized[time].iter().map(|v|v*c.ssm_in_multiplier).collect();
                let mut projected=linear(&input,in_proj);
                for (i,v) in projected.iter_mut().enumerate(){
                    let scale=if i<md{m[0]}else if i<2*md{m[1]}else if i<2*md+c.groups*s{m[2]}
                        else if i<2*md+2*c.groups*s{m[3]}else{m[4]}; *v*=scale;
                }
                let gate=&projected[..md]; let conv_input=&projected[md..md+cd];
                let dt=&projected[md+cd..];
                let mut conv_output=vec![0f32;cd];
                for channel in 0..cd {
                    conv_history[channel].rotate_left(1);
                    conv_history[channel][c.conv-1]=conv_input[channel];
                    let mut v=cb.data[channel];
                    // PyTorch Conv1d kernel orders oldest->newest.
                    for k in 0..c.conv{v+=cw.data[channel*c.conv+k]*conv_history[channel][k]}
                    conv_output[channel]=silu(v);
                }
                let xx=&conv_output[..md];
                let b=&conv_output[md..md+c.groups*s];
                let cc=&conv_output[md+c.groups*s..];
                let mut output=vec![0f32;md];
                for h in 0..c.mamba_heads {
                    let group=h/(c.mamba_heads/c.groups);
                    let step=softplus(dt[h]+dt_bias.data[h]);
                    let decay=(-a_log.data[h].exp()*step).exp();
                    for k in 0..hd {
                        let feature=h*hd+k;
                        let u=xx[feature];
                        let mut y=d_skip.data[h]*u;
                        for st in 0..s {
                            let i=feature*s+st;
                            recurrent[i]=decay*recurrent[i]+step*b[group*s+st]*u;
                            y+=cc[group*s+st]*recurrent[i];
                        }
                        output[feature]=y*silu(gate[feature]);
                    }
                }
                mamba_out[time]=linear(&output,out_proj).iter().map(|v|v*c.ssm_out_multiplier).collect();
            }
            // GQA: full-prefix causal softmax, per-query RoPE and K multiplier.
            let mut q=vec![vec![0f32;c.heads*c.head_dim];n];
            let mut k=vec![vec![0f32;c.kv_heads*c.head_dim];n];
            let mut v=vec![vec![0f32;c.kv_heads*c.head_dim];n];
            for t in 0..n {
                let a:Vec<f32>=normalized[t].iter().map(|x|x*c.attention_in_multiplier).collect();
                q[t]=linear(&a,self.t(&format!("{p}self_attn.q_proj.weight")));
                k[t]=linear(&a,self.t(&format!("{p}self_attn.k_proj.weight")));
                v[t]=linear(&a,self.t(&format!("{p}self_attn.v_proj.weight")));
                for head in 0..c.heads {rope(&mut q[t][head*c.head_dim..(head+1)*c.head_dim],t,c.rope_theta)}
                for head in 0..c.kv_heads {
                    let chunk=&mut k[t][head*c.head_dim..(head+1)*c.head_dim];
                    for x in chunk.iter_mut(){*x*=c.key_multiplier;}
                    rope(chunk,t,c.rope_theta);
                }
            }
            for t in 0..n {
                let mut y=vec![0f32;c.heads*c.head_dim];
                for h in 0..c.heads {
                    let kv=h/(c.heads/c.kv_heads);
                    let qi=&q[t][h*c.head_dim..(h+1)*c.head_dim];
                    let mut scores=Vec::with_capacity(t+1);
                    for j in 0..=t {
                        let kj=&k[j][kv*c.head_dim..(kv+1)*c.head_dim];
                        scores.push(qi.iter().zip(kj).map(|(a,b)|a*b).sum::<f32>()/(c.head_dim as f32).sqrt());
                    }
                    let mx=scores.iter().copied().fold(f32::NEG_INFINITY,f32::max);
                    let exps:Vec<f32>=scores.iter().map(|z|(z-mx).exp()).collect();
                    let denom=exps.iter().sum::<f32>();
                    for j in 0..=t {
                        let vv=&v[j][kv*c.head_dim..(kv+1)*c.head_dim];
                        for d in 0..c.head_dim{y[h*c.head_dim+d]+=exps[j]/denom*vv[d];}
                    }
                }
                let attn=linear(&y,self.t(&format!("{p}self_attn.o_proj.weight")));
                for i in 0..c.hidden {x[t][i]+=mamba_out[t][i]+c.attention_out_multiplier*attn[i];}
            }
            // pre-FF RMSNorm + SwiGLU with Falcon-H1 µP multipliers
            for t in 0..n {
                let a=norm(&x[t],self.t(&format!("{p}pre_ff_layernorm.weight")),c.eps);
                let up=linear(&a,self.t(&format!("{p}feed_forward.up_proj.weight")));
                let gate=linear(&a,self.t(&format!("{p}feed_forward.gate_proj.weight")));
                let z:Vec<f32>=up.iter().zip(&gate).map(|(u,g)|u*silu(g*c.mlp_multipliers[0])).collect();
                let ff=linear(&z,self.t(&format!("{p}feed_forward.down_proj.weight")));
                for i in 0..c.hidden{x[t][i]+=ff[i]*c.mlp_multipliers[1]}
            }
        }
        let final_h=norm(x.last().unwrap(),self.t("model.final_layernorm.weight"),c.eps);
        Ok(linear(&final_h,self.t("lm_head.weight")).iter().map(|z|z*c.lm_head_multiplier).collect())
    }
    pub fn generate(&self,ids:&[usize],count:usize)->Result<Vec<usize>,String>{
        if count>64 {return Err("max_new_tokens limited to 64 in CPU reference".into())}
        let mut sequence=ids.to_vec();
        for _ in 0..count {
            let y=self.logits(&sequence)?;
            let next=y.iter().enumerate().max_by(|a,b|a.1.total_cmp(b.1)).unwrap().0;
            sequence.push(next);
        }
        Ok(sequence)
    }
}

/// Deterministic tiny tensors for API unit tests. Not a pretrained model.
pub fn synthetic_weights(c:&Config)->Weights{
    let mut w=HashMap::new();let mut seed:u64=0x9e3779b97f4a7c15;
    for (name,shape) in expected_shapes(c){
        let mut data=vec![0f32;shape.iter().product()];
        for v in &mut data {seed ^= seed>>12;seed ^= seed<<25;seed ^= seed>>27;
            let z=(seed.wrapping_mul(2685821657736338717)>>40) as f32/(1u32<<24) as f32;
            *v=(z-0.5)*0.08;}
        if name.ends_with("layernorm.weight") || name.ends_with("final_layernorm.weight") {
            data.fill(1.);
        }
        if name.ends_with("mamba.D") {data.fill(1.);}
        if name.ends_with("mamba.dt_bias") {data.fill(0.);}
        if name.ends_with("mamba.A_log") {for (i,x) in data.iter_mut().enumerate(){*x=((i+1) as f32).ln()}}
        w.insert(name,Tensor{shape,data});
    }
    w
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test] fn test_falcon_dimensions(){assert!(Config::falcon_h1_05b().validate().is_ok());}
    #[test] fn test_cpu_logits(){let c=Config::demo();let m=Model::new(c.clone(),synthetic_weights(&c)).unwrap();let y=m.logits(&[1,2,3]).unwrap();assert_eq!(y.len(),c.vocab);assert!(y.iter().all(|x|x.is_finite()));}
    #[test] fn test_causal_prefix(){let c=Config::demo();let m=Model::new(c.clone(),synthetic_weights(&c)).unwrap();
        let a=m.logits(&[1,2]).unwrap();let b=m.logits(&[0,1,2]).unwrap();
        assert_eq!(a.len(),b.len()); // shifted positions legitimately alter RoPE
        assert!(m.logits(&[1,2,3]).is_ok());}
    #[test] fn test_determinism(){let c=Config::demo();let m=Model::new(c.clone(),synthetic_weights(&c)).unwrap();
        assert_eq!(m.logits(&[3,5,7]).unwrap(),m.logits(&[3,5,7]).unwrap());}
    #[test] fn test_generate(){let c=Config::demo();let m=Model::new(c.clone(),synthetic_weights(&c)).unwrap();
        assert_eq!(m.generate(&[1,2],2).unwrap().len(),4);}
    #[test] fn test_bad_token(){let c=Config::demo();let m=Model::new(c.clone(),synthetic_weights(&c)).unwrap();
        assert!(m.logits(&[c.vocab]).is_err());}
    #[test] fn test_weight_validation(){let c=Config::demo();let mut w=synthetic_weights(&c);
        w.remove("model.embed_tokens.weight");assert!(Model::new(c,w).is_err());}
}

#[cfg(feature = "candle")]
pub mod candle_backend;