package dibaguard

import (
	"fmt"
	"math"
)

// SeqNet is the tier-2 sequence network (dibaguard/flows/seqnet.py): event encoder -> causal GRU (one hidden
// state per source, one step per flow) -> site-conditioned head. Weights come from the flat export of seqnet.npz.
type SeqNet struct {
	EmbPort, EmbProto           []float32 // [PortVocab*16], [ProtoVocab*4]
	Enc0W, Enc0B, Enc1W, Enc1B  []float32 // [D*In], [D], [D*D], [D]
	WIH, WHH, BIH, BHH          [][]float32
	CtxW, CtxB                  []float32 // [2D*D], [2D] (nil without context)
	Head0W, Head0B              []float32 // [64*D], [64]
	Head1W, Head1B              []float32 // [C*64], [C]
	Gamma, Beta                 []float32 // [D] from the active site context
	D, In, Layers, C, NumIn     int
	PortVocab, ProtoVocab       int
	Temp                        float64
	Info                        map[string]any
}

// LoadSeqNet reads a flat export of SeqNet and applies the shipped default context when present.
func LoadSeqNet(path string) (*SeqNet, error) {
	f, err := LoadFlat(path)
	if err != nil {
		return nil, err
	}
	m := &SeqNet{Info: f.Info, Temp: 1}
	get := func(name string) ([]float32, []int) {
		if err != nil {
			return nil, nil
		}
		var a []float32
		var s []int
		a, s, err = f.Float32s(name)
		return a, s
	}
	var shp []int
	m.EmbPort, shp = get("emb_port")
	if shp != nil {
		m.PortVocab = shp[0]
	}
	m.EmbProto, shp = get("emb_proto")
	if shp != nil {
		m.ProtoVocab = shp[0]
	}
	m.Enc0W, shp = get("enc.0.w")
	if shp != nil {
		m.D, m.In = shp[0], shp[1]
		m.NumIn = m.In - 20
	}
	m.Enc0B, _ = get("enc.0.b")
	m.Enc1W, _ = get("enc.1.w")
	m.Enc1B, _ = get("enc.1.b")
	for l := 0; ; l++ {
		if _, ok := f.arrays[fmt.Sprintf("gru.l%d.w_hh", l)]; !ok {
			break
		}
		wih, _ := get(fmt.Sprintf("gru.l%d.w_ih", l))
		whh, _ := get(fmt.Sprintf("gru.l%d.w_hh", l))
		bih, _ := get(fmt.Sprintf("gru.l%d.b_ih", l))
		bhh, _ := get(fmt.Sprintf("gru.l%d.b_hh", l))
		m.WIH, m.WHH, m.BIH, m.BHH = append(m.WIH, wih), append(m.WHH, whh), append(m.BIH, bih), append(m.BHH, bhh)
		m.Layers++
	}
	m.Head0W, _ = get("head.0.w")
	m.Head0B, _ = get("head.0.b")
	m.Head1W, shp = get("head.1.w")
	if shp != nil {
		m.C = shp[0]
	}
	m.Head1B, _ = get("head.1.b")
	if _, ok := f.arrays["ctx.w"]; ok {
		m.CtxW, _ = get("ctx.w")
		m.CtxB, _ = get("ctx.b")
	}
	if err != nil {
		return nil, err
	}
	if v, ok := f.Info["temperature"].(float64); ok {
		m.Temp = v
	}
	m.Gamma, m.Beta = make([]float32, m.D), make([]float32, m.D)
	if _, ok := f.arrays["ctx_pool_default"]; ok {
		pool, _ := get("ctx_pool_default")
		m.SetContext(pool)
	}
	return m, nil
}

func gelu(x float32) float32 {
	y := float64(x)
	return float32(0.5 * y * (1 + math.Tanh(0.7978845608*(y+0.044715*y*y*y))))
}

func sigmoid32(x float32) float32 { return float32(1 / (1 + math.Exp(-float64(x)))) }

// SetContext derives (gamma, beta) from a pooled benign hidden vector [D] (the site profile). nil = none.
func (m *SeqNet) SetContext(pool []float32) {
	for j := range m.Gamma {
		m.Gamma[j], m.Beta[j] = 0, 0
	}
	if m.CtxW == nil || pool == nil {
		return
	}
	for i := 0; i < 2*m.D; i++ {
		s := m.CtxB[i]
		for j := 0; j < m.D; j++ {
			s += m.CtxW[i*m.D+j] * pool[j]
		}
		g := float32(math.Tanh(float64(s)))
		if i < m.D {
			m.Gamma[i] = g
		} else {
			m.Beta[i-m.D] = g
		}
	}
}

// NewState allocates one source's GRU state (Layers*D zeros).
func (m *SeqNet) NewState() []float32 { return make([]float32, m.Layers*m.D) }

// Step consumes one flow of a source: numeric features [NumIn], port id, proto id, the source's state (updated
// in place) -> logits [C] (index 0 = attack). scratch may be nil.
func (m *SeqNet) Step(x []float32, port, proto int, state []float32, out, scratch []float32) []float32 {
	D := m.D
	need := m.In + 2*D + 3*D + 3*D + 64
	if cap(scratch) < need {
		scratch = make([]float32, need)
	}
	scratch = scratch[:need]
	z := scratch[:m.In]
	copy(z, x[:m.NumIn])
	copy(z[m.NumIn:], m.EmbPort[port*16:port*16+16])
	copy(z[m.NumIn+16:], m.EmbProto[proto*4:proto*4+4])
	h1 := scratch[m.In : m.In+D]
	for j := 0; j < D; j++ {
		s := m.Enc0B[j]
		row := m.Enc0W[j*m.In : j*m.In+m.In]
		for i, zi := range z {
			s += row[i] * zi
		}
		h1[j] = gelu(s)
	}
	inp := scratch[m.In+D : m.In+2*D]
	for j := 0; j < D; j++ {
		s := m.Enc1B[j]
		row := m.Enc1W[j*D : j*D+D]
		for i, hi := range h1 {
			s += row[i] * hi
		}
		inp[j] = s
	}
	gi := scratch[m.In+2*D : m.In+5*D]
	gh := scratch[m.In+5*D : m.In+8*D]
	for l := 0; l < m.Layers; l++ {
		h := state[l*D : l*D+D]
		wih, whh := m.WIH[l], m.WHH[l]
		for j := 0; j < 3*D; j++ {
			s, t := m.BIH[l][j], m.BHH[l][j]
			ri, rh := wih[j*D:j*D+D], whh[j*D:j*D+D]
			for i := 0; i < D; i++ {
				s += ri[i] * inp[i]
				t += rh[i] * h[i]
			}
			gi[j], gh[j] = s, t
		}
		for j := 0; j < D; j++ {
			r := sigmoid32(gi[j] + gh[j])
			zg := sigmoid32(gi[D+j] + gh[D+j])
			n := float32(math.Tanh(float64(gi[2*D+j] + r*gh[2*D+j])))
			h[j] = (1-zg)*n + zg*h[j]
		}
		copy(inp, h)
	}
	hh := scratch[m.In+8*D : m.In+8*D+64]
	top := state[(m.Layers-1)*D : m.Layers*D]
	for j := 0; j < 64; j++ {
		s := m.Head0B[j]
		row := m.Head0W[j*D : j*D+D]
		for i := 0; i < D; i++ {
			s += row[i] * (top[i]*(1+m.Gamma[i]) + m.Beta[i])
		}
		hh[j] = gelu(s)
	}
	if cap(out) < m.C {
		out = make([]float32, m.C)
	}
	out = out[:m.C]
	for c := 0; c < m.C; c++ {
		s := m.Head1B[c]
		row := m.Head1W[c*64 : c*64+64]
		for i, hi := range hh {
			s += row[i] * hi
		}
		out[c] = s
	}
	return out
}

// ProbaAttack is sigmoid(logit0 / temperature).
func (m *SeqNet) ProbaAttack(logits []float32) float64 {
	return 1 / (1 + math.Exp(-float64(logits[0])/m.Temp))
}