package convert

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"io/fs"
	"maps"
	"math"
	"os"
	"slices"
	"strconv"
	"strings"

	"github.com/d4l3k/go-bfloat16"
	"github.com/pdevine/tensor"
	"github.com/pdevine/tensor/native"
	"github.com/x448/float16"

	"github.com/ollama/ollama/fs/ggml"
)

type qwen3NextRopeScaling struct {
	Type         string     `json:"type"`
	Factor       ropeFactor `json:"factor"`
	MropeSection []int32    `json:"mrope_section"`
}

type qwen3NextRopeParams struct {
	MRopeInterleaved    bool    `json:"mrope_interleaved"`
	MropeSection        []int32 `json:"mrope_section"`
	RopeType            string  `json:"rope_type"`
	RopeTheta           float32 `json:"rope_theta"`
	PartialRotaryFactor float32 `json:"partial_rotary_factor"`
}

type qwen3NextTextConfig struct {
	MaxPositionEmbeddings uint32  `json:"max_position_embeddings"`
	HiddenSize            uint32  `json:"hidden_size"`
	NumHiddenLayers       uint32  `json:"num_hidden_layers"`
	NumNextNPredictLayers uint32  `json:"num_nextn_predict_layers"`
	MTPNumHiddenLayers    uint32  `json:"mtp_num_hidden_layers"`
	IntermediateSize      uint32  `json:"intermediate_size"`
	NumAttentionHeads     uint32  `json:"num_attention_heads"`
	NumKeyValueHeads      uint32  `json:"num_key_value_heads"`
	HeadDim               uint32  `json:"head_dim"`
	RopeTheta             float32 `json:"rope_theta"`
	RMSNormEPS            float32 `json:"rms_norm_eps"`

	// MoE config
	NumExperts             uint32 `json:"num_experts"`
	NumExpertsPerToken     uint32 `json:"num_experts_per_tok"`
	NormTopkProb           *bool  `json:"norm_topk_prob"`
	MoEIntermediateSize    uint32 `json:"moe_intermediate_size"`
	SharedExpertIntermSize uint32 `json:"shared_expert_intermediate_size"`

	// Hybrid attention config
	FullAttentionInterval uint32   `json:"full_attention_interval"`
	LayerTypes            []string `json:"layer_types"`

	// Linear attention (Gated Delta Net) config
	LinearConvKernelDim uint32 `json:"linear_conv_kernel_dim"`
	LinearKeyHeadDim    uint32 `json:"linear_key_head_dim"`
	LinearNumKeyHeads   uint32 `json:"linear_num_key_heads"`
	LinearNumValueHeads uint32 `json:"linear_num_value_heads"`
	LinearValueHeadDim  uint32 `json:"linear_value_head_dim"`

	// RoPE config
	PartialRotaryFactor float32              `json:"partial_rotary_factor"`
	RopeScaling         qwen3NextRopeScaling `json:"rope_scaling"`
	RopeParameters      qwen3NextRopeParams  `json:"rope_parameters"`
}

type qwen3NextVisionConfig struct {
	Depth                  uint32  `json:"depth"`
	HiddenSize             uint32  `json:"hidden_size"`
	IntermediateSize       uint32  `json:"intermediate_size"`
	NumHeads               uint32  `json:"num_heads"`
	NumPositionEmbeddings  uint32  `json:"num_position_embeddings"`
	InChannels             uint32  `json:"in_channels"`
	OutHiddenSize          uint32  `json:"out_hidden_size"`
	PatchSize              uint32  `json:"patch_size"`
	SpatialMergeSize       uint32  `json:"spatial_merge_size"`
	RMSNormEps             float32 `json:"layer_norm_epsilon"`
	RopeTheta              float32 `json:"rope_theta"`
	TemporalPatchSize      uint32  `json:"temporal_patch_size"`
	DeepstackVisualIndexes []int32 `json:"deepstack_visual_indexes"`

	Size struct {
		ShortestEdge uint32 `json:"shortest_edge"`
		LongestEdge  uint32 `json:"longest_edge"`
	} `json:"size"`

	ImageMean []float32 `json:"image_mean"`
	ImageStd  []float32 `json:"image_std"`
}

type qwen3NextModel struct {
	ModelParameters
	qwen3NextTextConfig

	TextConfig  *qwen3NextTextConfig  `json:"text_config"`
	VisionModel qwen3NextVisionConfig `json:"vision_config"`

	ImageTokenID       uint32 `json:"image_token_id"`
	VisionStartTokenID uint32 `json:"vision_start_token_id"`
	VisionEndTokenID   uint32 `json:"vision_end_token_id"`
}

var (
	_ ModelConverter      = (*qwen3NextModel)(nil)
	_ MultimodalConverter = (*qwen3NextModel)(nil)
)

func (q *qwen3NextModel) parseMore(fsys fs.FS) error {
	if q.TextConfig != nil {
		q.qwen3NextTextConfig = *q.TextConfig
	}
	if q.NumNextNPredictLayers == 0 {
		q.NumNextNPredictLayers = q.MTPNumHiddenLayers
	}
	if q.NumNextNPredictLayers == 0 {
		nextn, err := qwen3NextInferNextNPredictLayers(fsys)
		if err != nil {
			return err
		}
		q.NumNextNPredictLayers = nextn
	}

	if q.RopeTheta == 0 {
		q.RopeTheta = q.RopeParameters.RopeTheta
	}
	if q.PartialRotaryFactor == 0 {
		q.PartialRotaryFactor = q.RopeParameters.PartialRotaryFactor
	}

	if q.RopeScaling.Type == "" && q.RopeParameters.RopeType != "" {
		q.RopeScaling.Type = q.RopeParameters.RopeType
	}

	// Pull vision preprocessing fields when present.
	if q.VisionModel.Depth > 0 {
		if bts, err := fs.ReadFile(fsys, "preprocessor_config.json"); err == nil {
			var pre struct {
				Size struct {
					ShortestEdge uint32 `json:"shortest_edge"`
					LongestEdge  uint32 `json:"longest_edge"`
				} `json:"size"`
				PatchSize         uint32    `json:"patch_size"`
				TemporalPatchSize uint32    `json:"temporal_patch_size"`
				MergeSize         uint32    `json:"merge_size"`
				ImageMean         []float32 `json:"image_mean"`
				ImageStd          []float32 `json:"image_std"`
			}
			if json.Unmarshal(bts, &pre) == nil {
				if q.VisionModel.PatchSize == 0 {
					q.VisionModel.PatchSize = pre.PatchSize
				}
				if q.VisionModel.TemporalPatchSize == 0 {
					q.VisionModel.TemporalPatchSize = pre.TemporalPatchSize
				}
				if q.VisionModel.SpatialMergeSize == 0 {
					q.VisionModel.SpatialMergeSize = pre.MergeSize
				}
				if q.VisionModel.Size.ShortestEdge == 0 {
					q.VisionModel.Size.ShortestEdge = pre.Size.ShortestEdge
				}
				if q.VisionModel.Size.LongestEdge == 0 {
					q.VisionModel.Size.LongestEdge = pre.Size.LongestEdge
				}
				if len(q.VisionModel.ImageMean) == 0 {
					q.VisionModel.ImageMean = pre.ImageMean
				}
				if len(q.VisionModel.ImageStd) == 0 {
					q.VisionModel.ImageStd = pre.ImageStd
				}
			}
		}
	}

	if q.NumHiddenLayers == 0 {
		return fmt.Errorf("qwen3next: num_hidden_layers must be set")
	}
	if q.NumAttentionHeads == 0 {
		return fmt.Errorf("qwen3next: num_attention_heads must be set")
	}
	if q.NumKeyValueHeads == 0 {
		return fmt.Errorf("qwen3next: num_key_value_heads must be set")
	}
	if q.HeadDim == 0 {
		return fmt.Errorf("qwen3next: head_dim must be set")
	}
	if q.RopeTheta == 0 {
		return fmt.Errorf("qwen3next: rope_theta must be set")
	}
	if q.PartialRotaryFactor <= 0 || q.PartialRotaryFactor > 1 {
		return fmt.Errorf("qwen3next: partial_rotary_factor must be in (0,1], got %v", q.PartialRotaryFactor)
	}
	if q.LinearNumKeyHeads == 0 || q.LinearNumValueHeads == 0 || q.LinearKeyHeadDim == 0 || q.LinearValueHeadDim == 0 {
		return fmt.Errorf("qwen3next: linear attention config must be set (linear_num_key_heads, linear_num_value_heads, linear_key_head_dim, linear_value_head_dim)")
	}
	if _, err := q.kvHeadCounts(); err != nil {
		return err
	}

	return nil
}

func qwen3NextInferNextNPredictLayers(fsys fs.FS) (uint32, error) {
	paths, err := fs.Glob(fsys, "*.safetensors")
	if err != nil {
		return 0, err
	}

	maxLayer := -1
	hasMTP := false
	for _, p := range paths {
		f, err := fsys.Open(p)
		if err != nil {
			return 0, err
		}

		var n int64
		if err := binary.Read(f, binary.LittleEndian, &n); err != nil {
			f.Close()
			return 0, err
		}

		b := bytes.NewBuffer(make([]byte, 0, n))
		if _, err = io.CopyN(b, f, n); err != nil {
			f.Close()
			return 0, err
		}
		f.Close()

		var headers map[string]safetensorMetadata
		if err := json.NewDecoder(b).Decode(&headers); err != nil {
			return 0, err
		}

		for name, value := range headers {
			if value.Type == "" || !strings.HasPrefix(name, "mtp.") {
				continue
			}
			hasMTP = true
			rest := strings.TrimPrefix(name, "mtp.layers.")
			layer, suffix, ok := strings.Cut(rest, ".")
			if !ok {
				continue
			}
			n, err := strconv.Atoi(layer)
			if err == nil && n > maxLayer && suffix != "" {
				maxLayer = n
			}
		}
	}

	if maxLayer >= 0 {
		return uint32(maxLayer + 1), nil
	}
	if hasMTP {
		return 1, nil
	}
	return 0, nil
}

func ConvertQwen35MTPDraft(fsys fs.FS, f *os.File, baseKV ggml.KV, baseTensors []*ggml.Tensor) error {
	arch := baseKV.Architecture()
	if arch != "qwen35" && arch != "qwen35moe" {
		return fmt.Errorf("MTP draft safetensors require a qwen3.5 base model, got %q", arch)
	}

	baseBlocks := baseKV.Uint("block_count")
	if baseBlocks == 0 {
		return fmt.Errorf("MTP draft safetensors require a base model with block_count")
	}
	if baseKV.Uint("nextn_predict_layers") > 0 {
		return fmt.Errorf("MTP draft safetensors require a base model without embedded MTP layers")
	}

	nextn, err := qwen3NextInferNextNPredictLayers(fsys)
	if err != nil {
		return err
	}
	if nextn == 0 {
		return fmt.Errorf("MTP draft safetensors did not contain mtp tensors")
	}

	q := &qwen3NextModel{
		qwen3NextTextConfig: qwen3NextTextConfig{
			NumHiddenLayers:       baseBlocks,
			NumNextNPredictLayers: nextn,
		},
	}
	ts, err := parseTensors(fsys, strings.NewReplacer(q.Replacements()...))
	if err != nil {
		return err
	}
	if err := ensureUniqueTensorNames(ts); err != nil {
		return err
	}

	mtpTensors := q.Tensors(ts)
	if len(mtpTensors) == 0 {
		return fmt.Errorf("MTP draft safetensors did not produce GGUF tensors")
	}
	for _, tensor := range mtpTensors {
		if !qwen35MTPDraftTensorName(tensor.Name, baseBlocks, nextn) {
			return fmt.Errorf("MTP draft safetensors produced unexpected tensor %q", tensor.Name)
		}
		tensor.Shape = slices.Clone(tensor.Shape)
		slices.Reverse(tensor.Shape)
	}

	kv := maps.Clone(baseKV)
	qwen35RemoveSplitMetadata(kv, arch)
	kv[arch+".block_count"] = baseBlocks + nextn
	kv[arch+".nextn_predict_layers"] = nextn

	tensors := make([]*ggml.Tensor, 0, len(baseTensors)+len(mtpTensors))
	tensors = append(tensors, baseTensors...)
	tensors = append(tensors, mtpTensors...)

	var parameters uint64
	for _, tensor := range tensors {
		parameters += tensor.Elements()
	}
	kv["general.parameter_count"] = parameters

	return ggml.WriteGGUF(f, kv, tensors)
}

func qwen35RemoveSplitMetadata(kv ggml.KV, arch string) {
	for _, key := range []string{
		"split.no",
		"split.count",
		"split.tensors.count",
	} {
		delete(kv, key)
		delete(kv, arch+"."+key)
	}
}

func qwen35MTPDraftTensorName(name string, base, nextn uint32) bool {
	for i := range nextn {
		if strings.HasPrefix(name, fmt.Sprintf("blk.%d.", base+i)) {
			return true
		}
	}
	return false
}

func (q *qwen3NextModel) kvHeadCounts() ([]uint32, error) {
	if len(q.LayerTypes) > 0 {
		kv := make([]uint32, q.NumHiddenLayers)
		hasFull := false
		hasRecurrent := false
		for i := range q.NumHiddenLayers {
			layerType := ""
			if i < uint32(len(q.LayerTypes)) {
				layerType = q.LayerTypes[i]
			}
			if layerType == "full_attention" {
				kv[i] = q.NumKeyValueHeads
				hasFull = true
			} else {
				hasRecurrent = true
			}
		}
		if !hasFull || !hasRecurrent {
			return nil, fmt.Errorf("qwen3next: layer_types must include both full_attention and linear_attention")
		}
		return kv, nil
	}

	if q.FullAttentionInterval == 0 {
		return nil, fmt.Errorf("qwen3next: full_attention_interval must be set")
	}
	if q.FullAttentionInterval > q.NumHiddenLayers {
		return nil, fmt.Errorf("qwen3next: full_attention_interval (%d) exceeds num_hidden_layers (%d)", q.FullAttentionInterval, q.NumHiddenLayers)
	}

	kv := make([]uint32, q.NumHiddenLayers)
	hasFull := false
	for i := range q.NumHiddenLayers {
		if (i+1)%q.FullAttentionInterval == 0 {
			kv[i] = q.NumKeyValueHeads
			hasFull = true
		}
	}
	if !hasFull {
		return nil, fmt.Errorf("qwen3next: head_count_kv would be all zeros (full_attention_interval=%d, num_hidden_layers=%d)", q.FullAttentionInterval, q.NumHiddenLayers)
	}
	return kv, nil
}

func (q *qwen3NextModel) ropeSections() []int32 {
	if len(q.RopeParameters.MropeSection) > 0 {
		return q.RopeParameters.MropeSection
	}
	return q.RopeScaling.MropeSection
}

func (q *qwen3NextModel) shouldReorderVHeads() bool {
	modelType := strings.ToLower(q.ModelType)
	if strings.Contains(modelType, "qwen3_next") || strings.Contains(modelType, "qwen3next") {
		return false
	}

	for _, arch := range q.Architectures {
		arch = strings.ToLower(arch)
		if strings.Contains(arch, "qwen3next") || strings.Contains(arch, "qwen3_next") {
			return false
		}
	}

	// Default to qwen3.5 layout for all other qwen3next-family imports.
	return true
}

func (q *qwen3NextModel) KV(t *Tokenizer) KV {
	kv := q.ModelParameters.KV(t)

	arch := "qwen35"
	if q.NumExperts > 0 {
		arch = "qwen35moe"
	}
	kv["general.architecture"] = arch
	kv["tokenizer.ggml.pre"] = "qwen35"
	kv["block_count"] = q.NumHiddenLayers + q.NumNextNPredictLayers
	if q.NumNextNPredictLayers > 0 {
		kv["nextn_predict_layers"] = q.NumNextNPredictLayers
	}
	kv["context_length"] = q.MaxPositionEmbeddings
	kv["embedding_length"] = q.HiddenSize
	kv["feed_forward_length"] = q.IntermediateSize
	kv["attention.head_count"] = q.NumAttentionHeads

	headDim := q.HeadDim
	if headDim == 0 && q.NumAttentionHeads > 0 {
		headDim = q.HiddenSize / q.NumAttentionHeads
	}
	kv["attention.key_length"] = headDim
	kv["attention.value_length"] = headDim
	kv["attention.layer_norm_rms_epsilon"] = q.RMSNormEPS
	kv["rope.freq_base"] = q.RopeTheta

	partialRotary := q.PartialRotaryFactor
	if partialRotary > 0 && partialRotary <= 1 {
		kv["rope.dimension_count"] = uint32(float32(headDim) * partialRotary)
	}

	if sections := q.ropeSections(); len(sections) > 0 {
		kv["mrope_sections"] = sections
		kv["rope.mrope_section"] = sections
		dimensionSections := append([]int32(nil), sections...)
		if len(dimensionSections) == 3 {
			dimensionSections = append(dimensionSections, 0)
		}
		kv["rope.dimension_sections"] = dimensionSections
	}
	if q.RopeParameters.MRopeInterleaved {
		kv["rope.mrope_interleaved"] = true
	}

	if q.RopeScaling.Type != "" && q.RopeScaling.Type != "default" {
		kv["rope.scaling.type"] = q.RopeScaling.Type
		kv["rope.scaling.factor"] = q.RopeScaling.Factor
	}

	if q.NumExperts > 0 {
		kv["expert_count"] = q.NumExperts
		kv["expert_used_count"] = q.NumExpertsPerToken
		if q.NormTopkProb != nil {
			kv["norm_top_k_prob"] = *q.NormTopkProb
		}
		if q.MoEIntermediateSize > 0 {
			kv["expert_feed_forward_length"] = q.MoEIntermediateSize
		}
		if q.SharedExpertIntermSize > 0 {
			kv["expert_shared_feed_forward_length"] = q.SharedExpertIntermSize
		}
	}

	dInner := q.LinearValueHeadDim * q.LinearNumValueHeads
	kv["ssm.inner_size"] = dInner
	kv["ssm.state_size"] = q.LinearKeyHeadDim
	kv["ssm.group_count"] = q.LinearNumKeyHeads
	kv["ssm.time_step_rank"] = q.LinearNumValueHeads
	kv["ssm.conv_kernel"] = q.LinearConvKernelDim
	if q.shouldReorderVHeads() {
		kv["ssm.v_head_reordered"] = true
	}
	if q.FullAttentionInterval > 0 {
		kv["full_attention_interval"] = q.FullAttentionInterval
	}

	if headCounts, err := q.kvHeadCounts(); err == nil {
		var maxKV uint32
		for _, count := range headCounts {
			if count > maxKV {
				maxKV = count
			}
		}
		kv["attention.head_count_kv"] = maxKV
	}

	if q.VisionModel.Depth > 0 {
		kv["vision.block_count"] = q.VisionModel.Depth
		kv["vision.embedding_length"] = q.VisionModel.HiddenSize
		if q.VisionModel.IntermediateSize > 0 {
			kv["vision.feed_forward_length"] = q.VisionModel.IntermediateSize
		}
		kv["vision.attention.head_count"] = q.VisionModel.NumHeads
		kv["vision.num_channels"] = q.VisionModel.InChannels
		if q.VisionModel.PatchSize > 0 {
			kv["vision.patch_size"] = q.VisionModel.PatchSize
		}
		if q.VisionModel.SpatialMergeSize > 0 {
			kv["vision.spatial_merge_size"] = q.VisionModel.SpatialMergeSize
		}
		if q.VisionModel.RMSNormEps > 0 {
			kv["vision.attention.layer_norm_epsilon"] = q.VisionModel.RMSNormEps
		}
		if q.VisionModel.RopeTheta > 0 {
			kv["vision.rope.freq_base"] = q.VisionModel.RopeTheta
		}
		if q.VisionModel.TemporalPatchSize > 0 {
			kv["vision.temporal_patch_size"] = q.VisionModel.TemporalPatchSize
		}
		kv["vision.deepstack_visual_indexes"] = q.VisionModel.DeepstackVisualIndexes
		if q.VisionModel.Size.ShortestEdge > 0 {
			kv["vision.shortest_edge"] = q.VisionModel.Size.ShortestEdge
		}
		if q.VisionModel.Size.LongestEdge > 0 {
			kv["vision.longest_edge"] = q.VisionModel.Size.LongestEdge
		}
		if len(q.VisionModel.ImageMean) > 0 {
			kv["vision.image_mean"] = q.VisionModel.ImageMean
		}
		if len(q.VisionModel.ImageStd) > 0 {
			kv["vision.image_std"] = q.VisionModel.ImageStd
		}
	}

	if q.ImageTokenID > 0 {
		kv["image_token_id"] = q.ImageTokenID
	}
	if q.VisionStartTokenID > 0 {
		kv["vision_start_token_id"] = q.VisionStartTokenID
	}
	if q.VisionEndTokenID > 0 {
		kv["vision_end_token_id"] = q.VisionEndTokenID
	}

	return kv
}

func (q *qwen3NextModel) TextKV(t *Tokenizer) KV {
	kv := q.KV(t)

	for _, key := range []string{
		"vision.block_count",
		"vision.embedding_length",
		"vision.feed_forward_length",
		"vision.attention.head_count",
		"vision.num_channels",
		"vision.patch_size",
		"vision.spatial_merge_size",
		"vision.attention.layer_norm_epsilon",
		"vision.rope.freq_base",
		"vision.temporal_patch_size",
		"vision.deepstack_visual_indexes",
		"vision.shortest_edge",
		"vision.longest_edge",
		"vision.image_mean",
		"vision.image_std",
		"image_token_id",
		"vision_start_token_id",
		"vision_end_token_id",
		"mrope_sections",
		"rope.mrope_section",
		"rope.mrope_interleaved",
		"ssm.v_head_reordered",
	} {
		delete(kv, key)
	}

	return kv
}

func (q *qwen3NextModel) ProjectorKV(*Tokenizer) KV {
	depth := q.VisionModel.Depth
	deepstack := make([]bool, depth)
	for _, idx := range q.VisionModel.DeepstackVisualIndexes {
		if idx >= 0 && uint32(idx) < depth {
			deepstack[idx] = true
		}
	}

	imageSize := uint32(768)
	if q.VisionModel.NumPositionEmbeddings > 0 && q.VisionModel.PatchSize > 0 {
		root := uint32(math.Sqrt(float64(q.VisionModel.NumPositionEmbeddings)))
		if root*root == q.VisionModel.NumPositionEmbeddings {
			imageSize = root * q.VisionModel.PatchSize
		}
	}

	projectionDim := q.VisionModel.OutHiddenSize
	if projectionDim == 0 {
		projectionDim = q.HiddenSize
	}
	layerNormEps := q.VisionModel.RMSNormEps
	if layerNormEps == 0 {
		layerNormEps = 1e-6
	}

	kv := KV{
		"general.architecture":                     "clip",
		"general.type":                             "mmproj",
		"general.file_type":                        uint32(1),
		"general.quantization_version":             uint32(2),
		"clip.has_vision_encoder":                  true,
		"clip.projector_type":                      "qwen3vl_merger",
		"clip.use_gelu":                            true,
		"clip.vision.block_count":                  depth,
		"clip.vision.embedding_length":             q.VisionModel.HiddenSize,
		"clip.vision.feed_forward_length":          q.VisionModel.IntermediateSize,
		"clip.vision.attention.head_count":         q.VisionModel.NumHeads,
		"clip.vision.image_size":                   imageSize,
		"clip.vision.patch_size":                   q.VisionModel.PatchSize,
		"clip.vision.projection_dim":               projectionDim,
		"clip.vision.spatial_merge_size":           q.VisionModel.SpatialMergeSize,
		"clip.vision.attention.layer_norm_epsilon": layerNormEps,
		"clip.vision.is_deepstack_layers":          deepstack,
	}
	if len(q.VisionModel.ImageMean) > 0 {
		kv["clip.vision.image_mean"] = q.VisionModel.ImageMean
	}
	if len(q.VisionModel.ImageStd) > 0 {
		kv["clip.vision.image_std"] = q.VisionModel.ImageStd
	}

	return kv
}

func (q *qwen3NextModel) TextTensors(ts []Tensor, _ *Tokenizer) []*ggml.Tensor {
	var text []Tensor
	for _, t := range ts {
		if qwen3NextVisionTensor(t.Name()) {
			continue
		}
		text = append(text, t)
	}

	return q.Tensors(text)
}

func (q *qwen3NextModel) ProjectorTensors(ts []Tensor) []*ggml.Tensor {
	if q.VisionModel.Depth == 0 {
		return nil
	}

	rename := strings.NewReplacer(
		"v.pos_embed", "v.position_embd",
		"v.patch_embed", "v.patch_embd",
		"v.merger.norm", "v.post_ln",
		"v.merger.linear_fc1", "mm.0",
		"v.merger.linear_fc2", "mm.2",
		".mlp.linear_fc1", ".ffn_up",
		".mlp.linear_fc2", ".ffn_down",
		".norm1", ".ln1",
		".norm2", ".ln2",
	)

	var out []*ggml.Tensor
	for _, t := range ts {
		name := t.Name()
		if !qwen3NextVisionTensor(name) {
			continue
		}

		if name == "v.patch_embed.weight" {
			out = append(out, q.qwen35PatchEmbedTensors(t)...)
			continue
		}

		outName := rename.Replace(name)
		kind := t.Kind()
		writer := io.WriterTo(t)
		if outName == "v.position_embd.weight" {
			kind = tensorKindFP32
			writer = tensorFloat32Writer{tensor: t}
		} else if sourceDType(t) == "BF16" && kind == tensorKindFP16 {
			kind = tensorKindBF16
			writer = tensorBF16Writer{tensor: t}
		}
		out = append(out, &ggml.Tensor{
			Name:     outName,
			Kind:     kind,
			Shape:    slices.Clone(t.Shape()),
			WriterTo: writer,
		})
	}

	return out
}

func qwen3NextVisionTensor(name string) bool {
	return strings.HasPrefix(name, "v.")
}

func (q *qwen3NextModel) qwen35PatchEmbedTensors(t Tensor) []*ggml.Tensor {
	shape := t.Shape()
	if len(shape) != 5 || shape[2] != 2 {
		return nil
	}

	outShape := []uint64{shape[0], shape[1], shape[3], shape[4]}
	return []*ggml.Tensor{
		{
			Name:     "v.patch_embd.weight",
			Kind:     tensorKindFP32,
			Shape:    slices.Clone(outShape),
			WriterTo: tensorFloat32Writer{tensor: t, repacker: q.qwen35PatchEmbedSlice(0)},
		},
		{
			Name:     "v.patch_embd.weight.1",
			Kind:     tensorKindFP32,
			Shape:    slices.Clone(outShape),
			WriterTo: tensorFloat32Writer{tensor: t, repacker: q.qwen35PatchEmbedSlice(1)},
		},
	}
}

func (q *qwen3NextModel) qwen35PatchEmbedSlice(slice int) Repacker {
	return func(_ string, data []float32, shape []uint64) ([]float32, error) {
		if len(shape) != 5 || shape[2] != 2 {
			return nil, fmt.Errorf("qwen3next: unexpected patch_embed shape %v", shape)
		}

		outChannels := int(shape[0])
		inChannels := int(shape[1])
		frames := int(shape[2])
		height := int(shape[3])
		width := int(shape[4])
		if slice < 0 || slice >= frames {
			return nil, fmt.Errorf("qwen3next: patch_embed slice %d out of range", slice)
		}

		expected := outChannels * inChannels * frames * height * width
		if len(data) != expected {
			return nil, fmt.Errorf("qwen3next: patch_embed data size %d, expected %d", len(data), expected)
		}

		out := make([]float32, outChannels*inChannels*height*width)
		for oc := range outChannels {
			for ic := range inChannels {
				for y := range height {
					for x := range width {
						src := ((((oc*inChannels+ic)*frames+slice)*height + y) * width) + x
						dst := (((oc*inChannels+ic)*height + y) * width) + x
						out[dst] = data[src]
					}
				}
			}
		}

		return out, nil
	}
}

type tensorBF16Writer struct {
	tensor   Tensor
	repacker Repacker
}

func (w tensorBF16Writer) WriteTo(dst io.Writer) (int64, error) {
	data, err := tensorFloat32Data(w.tensor)
	if err != nil {
		return 0, err
	}
	if w.repacker != nil {
		data, err = w.repacker(w.tensor.Name(), data, w.tensor.Shape())
		if err != nil {
			return 0, err
		}
	}

	u8s := bfloat16.EncodeFloat32(data)
	if _, err := dst.Write(u8s); err != nil {
		return 0, err
	}
	return int64(len(u8s)), nil
}

type tensorFloat32Writer struct {
	tensor   Tensor
	repacker Repacker
}

func (w tensorFloat32Writer) WriteTo(dst io.Writer) (int64, error) {
	data, err := tensorFloat32Data(w.tensor)
	if err != nil {
		return 0, err
	}
	if w.repacker != nil {
		data, err = w.repacker(w.tensor.Name(), data, w.tensor.Shape())
		if err != nil {
			return 0, err
		}
	}
	if err := binary.Write(dst, binary.LittleEndian, data); err != nil {
		return 0, err
	}
	return int64(len(data) * 4), nil
}

func tensorFloat32Data(t Tensor) ([]float32, error) {
	if st, ok := tensorSafetensor(t); ok {
		return safetensorFloat32Data(st)
	}

	var buf bytes.Buffer
	if _, err := t.WriteTo(&buf); err != nil {
		return nil, err
	}

	switch t.Kind() {
	case tensorKindFP32:
		out := make([]float32, buf.Len()/4)
		if err := binary.Read(bytes.NewReader(buf.Bytes()), binary.LittleEndian, out); err != nil {
			return nil, err
		}
		return out, nil
	case tensorKindFP16:
		raw := make([]uint16, buf.Len()/2)
		if err := binary.Read(bytes.NewReader(buf.Bytes()), binary.LittleEndian, raw); err != nil {
			return nil, err
		}
		out := make([]float32, len(raw))
		for i, v := range raw {
			out[i] = float16.Frombits(v).Float32()
		}
		return out, nil
	case tensorKindBF16:
		return bfloat16.DecodeFloat32(buf.Bytes()), nil
	default:
		return nil, fmt.Errorf("unsupported tensor kind %d for F32 writer", t.Kind())
	}
}

func tensorSafetensor(t Tensor) (safetensor, bool) {
	switch t := t.(type) {
	case safetensor:
		return t, true
	case *safetensor:
		return *t, true
	default:
		return safetensor{}, false
	}
}

func safetensorFloat32Data(st safetensor) ([]float32, error) {
	f, err := st.fs.Open(st.path)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	var r io.Reader
	if readerAt, ok := f.(io.ReaderAt); ok {
		r = io.NewSectionReader(readerAt, st.offset, st.size)
	} else if seeker, ok := f.(io.Seeker); ok {
		if _, err := seeker.Seek(st.offset, io.SeekStart); err != nil {
			return nil, err
		}
		r = f
	} else {
		if _, err := io.CopyN(io.Discard, f, st.offset); err != nil {
			return nil, err
		}
		r = f
	}

	br := bufio.NewReaderSize(r, min(32<<10, int(st.size)))
	var out []float32
	switch st.dtype {
	case "F32":
		out = make([]float32, st.size/4)
		if err := binary.Read(br, binary.LittleEndian, out); err != nil {
			return nil, err
		}
	case "F16":
		raw := make([]uint16, st.size/2)
		if err := binary.Read(br, binary.LittleEndian, raw); err != nil {
			return nil, err
		}
		out = make([]float32, len(raw))
		for i, v := range raw {
			out[i] = float16.Frombits(v).Float32()
		}
	case "BF16":
		raw := make([]uint8, st.size)
		if err := binary.Read(br, binary.LittleEndian, raw); err != nil {
			return nil, err
		}
		out = bfloat16.DecodeFloat32(raw)
	case "F8_E4M3":
		raw := make([]uint8, st.size)
		if err := binary.Read(br, binary.LittleEndian, raw); err != nil {
			return nil, err
		}
		out, err = st.decodeFP8E4M3(raw)
		if err != nil {
			return nil, err
		}
	default:
		return nil, fmt.Errorf("unsupported safetensor dtype %q", st.dtype)
	}

	if st.repacker != nil {
		out, err = st.repacker(st.Name(), out, st.Shape())
		if err != nil {
			return nil, err
		}
	}
	return out, nil
}

func (q *qwen3NextModel) Tensors(ts []Tensor) []*ggml.Tensor {
	var out []*ggml.Tensor

	ts = q.renameMTPLayerTensors(ts)

	blockCount := q.NumHiddenLayers + q.NumNextNPredictLayers
	merges := make([]merge, blockCount*3)
	for i := range blockCount {
		merges[i*3+0] = merge{
			fmt.Sprintf("blk.%d.mlp.experts.*.gate_proj.weight", i),
			fmt.Sprintf("blk.%d.ffn_gate_exps.weight", i),
		}
		merges[i*3+1] = merge{
			fmt.Sprintf("blk.%d.mlp.experts.*.up_proj.weight", i),
			fmt.Sprintf("blk.%d.ffn_up_exps.weight", i),
		}
		merges[i*3+2] = merge{
			fmt.Sprintf("blk.%d.mlp.experts.*.down_proj.weight", i),
			fmt.Sprintf("blk.%d.ffn_down_exps.weight", i),
		}
	}

	merged, remaining := mergeTensors(ts, merges...)
	out = append(out, merged...)

	for _, t := range remaining {
		name := t.Name()
		shape := t.Shape()

		if names := q.mtpTensorNames(name); len(names) > 0 {
			for _, name := range names {
				out = q.appendDirectTensor(out, t, name)
			}
			continue
		}

		if strings.HasSuffix(name, ".ssm_in.weight") {
			if qkv, gate, ok := q.splitQKVZTensor(t); ok {
				out = append(out, qkv, gate)
				continue
			}
			panic(fmt.Sprintf("qwen3next: failed to split %s into attn_qkv/attn_gate (shape=%v)", name, shape))
		}

		switch {
		case strings.Contains(name, ".mlp.experts.gate_up_proj"):
			out = append(out, slices.Collect(splitDim(t, 1,
				split{Replacer: strings.NewReplacer(".mlp.experts.gate_up_proj", ".ffn_gate_exps.weight")},
				split{Replacer: strings.NewReplacer(".mlp.experts.gate_up_proj", ".ffn_up_exps.weight")},
			))...)

		case strings.Contains(name, ".mlp.experts.down_proj"):
			out = append(out, &ggml.Tensor{
				Name:     strings.NewReplacer(".mlp.experts.down_proj", ".ffn_down_exps.weight").Replace(name),
				Kind:     t.Kind(),
				Shape:    slices.Clone(shape),
				WriterTo: t,
			})

		case strings.HasPrefix(name, "v.blk.") && strings.Contains(name, ".attn_qkv"):
			out = append(out, slices.Collect(splitDim(t, 0,
				split{Replacer: strings.NewReplacer("attn_qkv", "attn_q")},
				split{Replacer: strings.NewReplacer("attn_qkv", "attn_k")},
				split{Replacer: strings.NewReplacer("attn_qkv", "attn_v")},
			))...)

		case strings.Contains(name, "patch_embed") && strings.HasSuffix(name, "weight"):
			out = append(out, &ggml.Tensor{
				Name:     name,
				Kind:     t.Kind(),
				Shape:    append([]uint64{shape[0] * shape[1]}, shape[2:]...),
				WriterTo: t,
			})

		case strings.HasSuffix(name, "_norm.weight") && !strings.HasSuffix(name, ".ssm_norm.weight"):
			t.SetRepacker(q.addOne)
			out = append(out, &ggml.Tensor{Name: name, Kind: t.Kind(), Shape: slices.Clone(shape), WriterTo: t})

		case strings.HasSuffix(name, ".ssm_a"):
			t.SetRepacker(q.repackSSMA())
			out = append(out, &ggml.Tensor{Name: name, Kind: t.Kind(), Shape: slices.Clone(shape), WriterTo: t})

		case strings.HasSuffix(name, ".attn_qkv.weight"):
			if q.shouldReorderVHeads() {
				t.SetRepacker(q.repackAttnQKV())
			}
			out = append(out, &ggml.Tensor{Name: name, Kind: t.Kind(), Shape: slices.Clone(shape), WriterTo: t})

		case strings.HasSuffix(name, ".attn_gate.weight"):
			if q.shouldReorderVHeads() {
				// HF tensor layout is [out_features, in_features]; reorder rows.
				t.SetRepacker(q.repackReorderDim(0, int(q.LinearValueHeadDim)))
			}
			out = append(out, &ggml.Tensor{Name: name, Kind: t.Kind(), Shape: slices.Clone(shape), WriterTo: t})

		case strings.HasSuffix(name, ".ssm_beta.weight"), strings.HasSuffix(name, ".ssm_alpha.weight"):
			if q.shouldReorderVHeads() {
				// HF tensor layout is [out_features, in_features]; reorder rows.
				t.SetRepacker(q.repackReorderDim(0, 1))
			}
			out = append(out, &ggml.Tensor{Name: name, Kind: t.Kind(), Shape: slices.Clone(shape), WriterTo: t})

		case strings.HasSuffix(name, ".ssm_dt"), strings.HasSuffix(name, ".ssm_dt.bias"):
			if q.shouldReorderVHeads() {
				t.SetRepacker(q.repackReorderDim(0, 1))
			}
			out = append(out, &ggml.Tensor{Name: name, Kind: t.Kind(), Shape: slices.Clone(shape), WriterTo: t})

		case strings.HasSuffix(name, ".ssm_out.weight"):
			if q.shouldReorderVHeads() {
				// HF out_proj layout is [out_features, in_features]; reorder columns.
				t.SetRepacker(q.repackReorderDim(1, int(q.LinearValueHeadDim)))
			}
			out = append(out, &ggml.Tensor{Name: name, Kind: t.Kind(), Shape: slices.Clone(shape), WriterTo: t})

		case strings.HasSuffix(name, ".ssm_conv1d.weight"):
			newShape := slices.Clone(shape)
			if len(shape) == 3 {
				if shape[0] == 1 {
					newShape = []uint64{shape[1], shape[2]}
				} else if shape[1] == 1 {
					newShape = []uint64{shape[0], shape[2]}
				}
			}
			if q.shouldReorderVHeads() {
				t.SetRepacker(q.repackConv1D())
			}
			out = append(out, &ggml.Tensor{Name: name, Kind: t.Kind(), Shape: newShape, WriterTo: t})

		default:
			out = append(out, &ggml.Tensor{Name: name, Kind: t.Kind(), Shape: slices.Clone(shape), WriterTo: t})
		}
	}

	return out
}

func (q *qwen3NextModel) renameMTPLayerTensors(ts []Tensor) []Tensor {
	var out []Tensor
	for i, t := range ts {
		name, ok := q.mtpLayerTensorName(t.Name())
		if !ok {
			continue
		}
		if out == nil {
			out = slices.Clone(ts)
		}
		out[i] = &renamedTensor{Tensor: t, name: name}
	}
	if out != nil {
		return out
	}
	return ts
}

func (q *qwen3NextModel) mtpLayerTensorName(name string) (string, bool) {
	rest := strings.TrimPrefix(name, "mtp.layers.")
	if rest == name {
		return "", false
	}
	layer, suffix, ok := strings.Cut(rest, ".")
	if !ok {
		return "", false
	}
	idx, err := strconv.ParseUint(layer, 10, 32)
	if err != nil {
		return "", false
	}
	return fmt.Sprintf("blk.%d.%s", q.NumHiddenLayers+uint32(idx), suffix), true
}

type renamedTensor struct {
	Tensor
	name string
}

func (t *renamedTensor) Name() string {
	return t.name
}

func (t *renamedTensor) Clone() Tensor {
	return &renamedTensor{Tensor: t.Tensor.Clone(), name: t.name}
}

func (t *renamedTensor) SourceDType() string {
	return sourceDType(t.Tensor)
}

func (q *qwen3NextModel) appendDirectTensor(out []*ggml.Tensor, t Tensor, name string) []*ggml.Tensor {
	if qwen3NextShouldShiftNorm(name) {
		t = t.Clone()
		t.SetRepacker(q.addOne)
	}
	return append(out, &ggml.Tensor{Name: name, Kind: t.Kind(), Shape: slices.Clone(t.Shape()), WriterTo: t})
}

func qwen3NextShouldShiftNorm(name string) bool {
	if strings.HasSuffix(name, ".ssm_norm.weight") {
		return false
	}
	return strings.HasSuffix(name, "_norm.weight") ||
		strings.HasSuffix(name, ".nextn.enorm.weight") ||
		strings.HasSuffix(name, ".nextn.hnorm.weight")
}

func (q *qwen3NextModel) mtpTensorNames(name string) []string {
	if !strings.HasPrefix(name, "mtp.") {
		return nil
	}

	base := q.NumHiddenLayers
	nextn := q.NumNextNPredictLayers
	if nextn == 0 {
		nextn = 1
	}

	var suffix string
	switch name {
	case "mtp.fc.weight":
		suffix = "nextn.eh_proj.weight"
	case "mtp.pre_fc_norm_embedding.weight":
		suffix = "nextn.enorm.weight"
	case "mtp.pre_fc_norm_hidden.weight":
		suffix = "nextn.hnorm.weight"
	case "mtp.norm.weight":
		suffix = "nextn.shared_head_norm.weight"
	case "mtp.embed_tokens.weight":
		suffix = "nextn.embed_tokens.weight"
	case "mtp.shared_head.head.weight":
		suffix = "nextn.shared_head_head.weight"
	case "mtp.shared_head.norm.weight":
		suffix = "nextn.shared_head_norm.weight"
	default:
		return nil
	}

	names := make([]string, 0, nextn)
	for i := range nextn {
		names = append(names, fmt.Sprintf("blk.%d.%s", base+i, suffix))
	}
	return names
}

func (q *qwen3NextModel) repackReorderDim(dim, headDim int) Repacker {
	return func(_ string, data []float32, shape []uint64) ([]float32, error) {
		if !q.shouldReorderVHeads() {
			return data, nil
		}
		numK := int(q.LinearNumKeyHeads)
		numVPerK := int(q.LinearNumValueHeads / q.LinearNumKeyHeads)
		return reorderHeadLayout(data, shape, dim, numK, numVPerK, headDim)
	}
}

func (q *qwen3NextModel) repackAttnQKV() Repacker {
	return func(_ string, data []float32, shape []uint64) ([]float32, error) {
		if !q.shouldReorderVHeads() || len(shape) != 2 {
			return data, nil
		}

		rows := int(shape[0])
		cols := int(shape[1])
		numK := int(q.LinearNumKeyHeads)
		numV := int(q.LinearNumValueHeads)
		headK := int(q.LinearKeyHeadDim)
		headV := int(q.LinearValueHeadDim)
		qDim := headK * numK
		kDim := headK * numK
		vDim := headV * numV
		qkvDim := qDim + kDim + vDim

		switch {
		case rows == qkvDim:
			// HF layout: [out_features, in_features]. Keep Q/K rows unchanged and
			// reorder only V rows from grouped -> tiled head layout.
			out := make([]float32, len(data))
			qkRows := qDim + kDim
			qkSize := qkRows * cols
			copy(out[:qkSize], data[:qkSize])

			vStart := qkSize
			vEnd := vStart + vDim*cols
			reorderedV, err := reorderHeadLayout(data[vStart:vEnd], []uint64{uint64(vDim), uint64(cols)}, 0, numK, numV/numK, headV)
			if err != nil {
				return nil, err
			}
			copy(out[vStart:vEnd], reorderedV)
			copy(out[vEnd:], data[vEnd:])
			return out, nil

		case cols == qkvDim:
			// Fallback for already-transposed [in_features, out_features] tensors.
			out := make([]float32, len(data))
			copy(out, data)
			for r := range rows {
				base := r * cols
				vStart := base + qDim + kDim
				vEnd := vStart + vDim
				reorderedV, err := reorderHeadLayout(out[vStart:vEnd], []uint64{uint64(vDim)}, 0, numK, numV/numK, headV)
				if err != nil {
					return nil, err
				}
				copy(out[vStart:vEnd], reorderedV)
			}
			return out, nil

		default:
			return data, nil
		}
	}
}

func (q *qwen3NextModel) repackConv1D() Repacker {
	return func(_ string, data []float32, shape []uint64) ([]float32, error) {
		if !q.shouldReorderVHeads() {
			return data, nil
		}

		normShape := slices.Clone(shape)
		if len(shape) == 3 {
			if shape[0] == 1 {
				normShape = []uint64{shape[1], shape[2]}
			} else if shape[1] == 1 {
				normShape = []uint64{shape[0], shape[2]}
			}
		}
		if len(normShape) != 2 {
			return data, nil
		}

		rows := int(normShape[0])
		cols := int(normShape[1])
		numK := int(q.LinearNumKeyHeads)
		numV := int(q.LinearNumValueHeads)
		headK := int(q.LinearKeyHeadDim)
		headV := int(q.LinearValueHeadDim)
		qkChannels := 2 * headK * numK
		totalChannels := qkChannels + headV*numV
		if qkChannels <= 0 {
			return data, nil
		}

		switch {
		case rows == totalChannels:
			// HF layout after squeeze: [channels, kernel]
			out := make([]float32, len(data))
			prefix := qkChannels * cols
			copy(out[:prefix], data[:prefix])
			reorderedV, err := reorderHeadLayout(data[prefix:], []uint64{uint64(totalChannels - qkChannels), uint64(cols)}, 0, numK, numV/numK, headV)
			if err != nil {
				return nil, err
			}
			copy(out[prefix:], reorderedV)
			return out, nil
		case cols == totalChannels:
			// Fallback for transposed [kernel, channels]
			out := make([]float32, len(data))
			copy(out, data)
			vChannels := totalChannels - qkChannels
			for r := range rows {
				base := r * cols
				vStart := base + qkChannels
				vEnd := vStart + vChannels
				reorderedV, err := reorderHeadLayout(out[vStart:vEnd], []uint64{uint64(vChannels)}, 0, numK, numV/numK, headV)
				if err != nil {
					return nil, err
				}
				copy(out[vStart:vEnd], reorderedV)
			}
			return out, nil
		default:
			return data, nil
		}
	}
}

func (q *qwen3NextModel) repackSSMA() Repacker {
	return func(_ string, data []float32, shape []uint64) ([]float32, error) {
		result := make([]float32, len(data))
		for i, v := range data {
			result[i] = -float32(math.Exp(float64(v)))
		}
		if !q.shouldReorderVHeads() {
			return result, nil
		}
		numK := int(q.LinearNumKeyHeads)
		numVPerK := int(q.LinearNumValueHeads / q.LinearNumKeyHeads)
		return reorderHeadLayout(result, shape, 0, numK, numVPerK, 1)
	}
}

func reorderHeadLayout(data []float32, shape []uint64, dim int, numKHeads, numVPerK, headDim int) ([]float32, error) {
	if len(shape) == 0 || numKHeads <= 0 || numVPerK <= 0 || headDim <= 0 {
		return data, nil
	}

	dims := make([]int, len(shape))
	for i := range shape {
		dims[i] = int(shape[i])
	}
	if dim < 0 {
		dim += len(dims)
	}
	if dim < 0 || dim >= len(dims) {
		return data, nil
	}

	expected := numKHeads * numVPerK * headDim
	if dims[dim] != expected {
		return data, nil
	}

	newShape := make([]int, 0, len(dims)+2)
	newShape = append(newShape, dims[:dim]...)
	newShape = append(newShape, numKHeads, numVPerK, headDim)
	newShape = append(newShape, dims[dim+1:]...)

	var tt tensor.Tensor = tensor.New(tensor.WithShape(dims...), tensor.WithBacking(data))
	if err := tt.Reshape(newShape...); err != nil {
		return nil, err
	}

	perm := make([]int, len(newShape))
	for i := range perm {
		perm[i] = i
	}
	perm[dim], perm[dim+1] = perm[dim+1], perm[dim]

	tt, err := tensor.Transpose(tt, perm...)
	if err != nil {
		return nil, err
	}
	tt = tensor.Materialize(tt)

	total := 1
	for _, d := range dims {
		total *= d
	}
	if err := tt.Reshape(total); err != nil {
		return nil, err
	}
	return native.VectorF32(tt.(*tensor.Dense))
}

type qkvzSplitSpec struct {
	hidden    int
	headKDim  int
	headVDim  int
	numKHeads int
	numVHeads int
	qkvzDim   int
	qkvOut    int
	gateOut   int
}

func (q *qwen3NextModel) qkvzSpec(shape []uint64) (qkvzSplitSpec, bool) {
	if len(shape) != 2 {
		return qkvzSplitSpec{}, false
	}

	numKHeads := int(q.LinearNumKeyHeads)
	numVHeads := int(q.LinearNumValueHeads)
	headKDim := int(q.LinearKeyHeadDim)
	headVDim := int(q.LinearValueHeadDim)
	if numKHeads == 0 || numVHeads == 0 || headKDim == 0 || headVDim == 0 {
		return qkvzSplitSpec{}, false
	}
	if numVHeads%numKHeads != 0 {
		return qkvzSplitSpec{}, false
	}

	hidden := int(shape[1])
	vPerHead := headVDim * (numVHeads / numKHeads)
	qkvzDim := 2*headKDim + 2*vPerHead
	expectedOut := qkvzDim * numKHeads
	if int(shape[0]) != expectedOut {
		return qkvzSplitSpec{}, false
	}

	return qkvzSplitSpec{
		hidden:    hidden,
		headKDim:  headKDim,
		headVDim:  headVDim,
		numKHeads: numKHeads,
		numVHeads: numVHeads,
		qkvzDim:   qkvzDim,
		qkvOut:    2*headKDim*numKHeads + headVDim*numVHeads,
		gateOut:   headVDim * numVHeads,
	}, true
}

func (q *qwen3NextModel) splitQKVZTensor(t Tensor) (*ggml.Tensor, *ggml.Tensor, bool) {
	spec, ok := q.qkvzSpec(t.Shape())
	if !ok {
		return nil, nil, false
	}

	qkvTensor := t.Clone()
	qkvTensor.SetRepacker(q.repackQKVZ(spec, false))

	gateTensor := t.Clone()
	gateTensor.SetRepacker(q.repackQKVZ(spec, true))

	qkvName := strings.Replace(t.Name(), "ssm_in", "attn_qkv", 1)
	gateName := strings.Replace(t.Name(), "ssm_in", "attn_gate", 1)

	return &ggml.Tensor{
			Name:     qkvName,
			Kind:     t.Kind(),
			Shape:    []uint64{uint64(spec.qkvOut), uint64(spec.hidden)},
			WriterTo: qkvTensor,
		}, &ggml.Tensor{
			Name:     gateName,
			Kind:     t.Kind(),
			Shape:    []uint64{uint64(spec.gateOut), uint64(spec.hidden)},
			WriterTo: gateTensor,
		}, true
}

func (q *qwen3NextModel) repackQKVZ(spec qkvzSplitSpec, extractGate bool) Repacker {
	vPerHead := spec.headVDim * (spec.numVHeads / spec.numKHeads)

	return func(_ string, data []float32, shape []uint64) ([]float32, error) {
		dims := make([]int, len(shape))
		for i := range shape {
			dims[i] = int(shape[i])
		}

		var tt tensor.Tensor = tensor.New(tensor.WithShape(dims...), tensor.WithBacking(data))
		var err error

		tt, err = tensor.Transpose(tt, 1, 0)
		if err != nil {
			return nil, err
		}
		tt = tensor.Materialize(tt)

		if err := tt.Reshape(spec.hidden, spec.numKHeads, spec.qkvzDim); err != nil {
			return nil, err
		}

		offset := 0
		qSlice, err := tt.Slice(nil, nil, tensor.S(offset, offset+spec.headKDim))
		if err != nil {
			return nil, err
		}
		offset += spec.headKDim
		kSlice, err := tt.Slice(nil, nil, tensor.S(offset, offset+spec.headKDim))
		if err != nil {
			return nil, err
		}
		offset += spec.headKDim
		vSlice, err := tt.Slice(nil, nil, tensor.S(offset, offset+vPerHead))
		if err != nil {
			return nil, err
		}
		offset += vPerHead
		zSlice, err := tt.Slice(nil, nil, tensor.S(offset, offset+vPerHead))
		if err != nil {
			return nil, err
		}

		qMat := tensor.Materialize(qSlice).(*tensor.Dense)
		kMat := tensor.Materialize(kSlice).(*tensor.Dense)
		vMat := tensor.Materialize(vSlice).(*tensor.Dense)
		zMat := tensor.Materialize(zSlice).(*tensor.Dense)

		if err := qMat.Reshape(spec.hidden, spec.numKHeads*spec.headKDim); err != nil {
			return nil, err
		}
		if err := kMat.Reshape(spec.hidden, spec.numKHeads*spec.headKDim); err != nil {
			return nil, err
		}
		if err := vMat.Reshape(spec.hidden, spec.numKHeads*vPerHead); err != nil {
			return nil, err
		}
		if err := zMat.Reshape(spec.hidden, spec.numKHeads*vPerHead); err != nil {
			return nil, err
		}

		var out tensor.Tensor
		if extractGate {
			out = zMat
		} else {
			out, err = tensor.Concat(1, qMat, kMat, vMat)
			if err != nil {
				return nil, err
			}
		}

		out = tensor.Materialize(out)
		out, err = tensor.Transpose(out, 1, 0)
		if err != nil {
			return nil, err
		}
		out = tensor.Materialize(out)

		if err := out.Reshape(out.Shape().TotalSize()); err != nil {
			return nil, err
		}

		return native.VectorF32(out.(*tensor.Dense))
	}
}

func (*qwen3NextModel) addOne(_ string, data []float32, shape []uint64) ([]float32, error) {
	n := tensor.New(tensor.WithShape(int(shape[0])), tensor.WithBacking(data))
	ones := tensor.Ones(tensor.Float32, int(shape[0]))

	n, err := n.Add(ones)
	if err != nil {
		return nil, err
	}

	ts, err := native.SelectF32(n, 0)
	if err != nil {
		return nil, err
	}

	var f32s []float32
	for _, t := range ts {
		f32s = append(f32s, t...)
	}

	return f32s, nil
}

func (q *qwen3NextModel) Replacements() []string {
	return []string{
		// Embeddings and output
		"lm_head", "output",
		"model.language_model.embed_tokens", "token_embd",
		"model.language_model.norm", "output_norm",
		"model.language_model.layers", "blk",
		"model.embed_tokens", "token_embd",
		"model.norm", "output_norm",
		"model.layers", "blk",

		// Vision
		"model.visual", "v",
		"patch_embed.proj", "patch_embed",
		"blocks", "blk",
		"attn.qkv", "attn_qkv",
		"attn.proj", "attn_out",
		"deepstack_merger_list", "deepstack_merger",

		// Layer norms
		"input_layernorm", "attn_norm",
		"post_attention_layernorm", "post_attention_norm",

		// Full attention (self_attn)
		"self_attn.q_proj", "attn_q",
		"self_attn.q_norm", "attn_q_norm",
		"self_attn.k_proj", "attn_k",
		"self_attn.k_norm", "attn_k_norm",
		"self_attn.v_proj", "attn_v",
		"self_attn.o_proj", "attn_output",

		// Linear attention (legacy qwen3next)
		"linear_attn.in_proj_qkvz", "ssm_in",
		"linear_attn.in_proj_ba", "ssm_ba",

		// Linear attention (qwen35)
		"linear_attn.in_proj_qkv", "attn_qkv",
		"linear_attn.in_proj_z", "attn_gate",
		"linear_attn.in_proj_a", "ssm_alpha",
		"linear_attn.in_proj_b", "ssm_beta",

		"linear_attn.conv1d", "ssm_conv1d",
		"linear_attn.dt_bias", "ssm_dt.bias",
		"linear_attn.dt_proj", "ssm_dt",
		"linear_attn.A_log", "ssm_a",
		"linear_attn.norm", "ssm_norm",
		"linear_attn.out_proj", "ssm_out",

		// MoE
		"mlp.gate.weight", "ffn_gate_inp.weight",
		"mlp.shared_expert.down_proj", "ffn_down_shexp",
		"mlp.shared_expert.gate_proj", "ffn_gate_shexp",
		"mlp.shared_expert.up_proj", "ffn_up_shexp",
		"mlp.shared_expert_gate", "ffn_gate_inp_shexp",

		// Dense FFN
		"mlp.down_proj", "ffn_down",
		"mlp.gate_proj", "ffn_gate",
		"mlp.up_proj", "ffn_up",
	}
}