import numpy as np
import math
import os
from dataclasses import dataclass
import yaml

@dataclass
class ModelConfig:
    sequence_length: int
    d_model: int
    num_heads: int
    num_encoders: int
    ff_dim: int

@dataclass
class KernelConfig:
    store_fp16: bool
    acc_fp16: bool
    encoder_layer_count: int
    

@dataclass
class Configuration:
    model: ModelConfig
    kernel: KernelConfig


def load_model_document(data_dir):
    """Read the model.yaml document that describes model and kernel dimensions."""
    yaml_path = f"{data_dir}model.yaml"
    if os.path.exists(yaml_path):
        with open(yaml_path, "r") as f:
            return yaml.safe_load(f) or {}

    raise FileNotFoundError(f"No model.yaml found in {data_dir}")

def cpp_float_repr(x):
    if math.isnan(x):
        return 'NAN'
    elif math.isinf(x):
        return 'INFINITY' if x > 0 else '-INFINITY'
    else:
        return f'{x:.6f}f'  # Append 'f' to represent a float in C++

def save_cpp_array_3d(name, arr, dtype='float'):
    M, N, K = arr.shape
    lines = [f'static const {dtype} {name}[{M}][{N}][{K}] = {{']
    for i in range(M):
        lines.append("    {")
        for j in range(N):
            row_vals = ", ".join(cpp_float_repr(v) for v in arr[i, j])
            lines.append(f"        {{ {row_vals} }},")
        lines.append("    },")
    lines.append("};")
    result = "\n".join(lines)

    os.makedirs("matrices", exist_ok=True)
    with open(f"matrices/{name}.h", "w") as f:
        f.write(result)

def save_cpp_array_2d(name, arr, dtype='float'):
    M, N = arr.shape
    lines = [f'static const {dtype} {name}[{M}][{N}] = {{']
    for i in range(M):
        row_vals = ", ".join(cpp_float_repr(v) for v in arr[i])
        lines.append(f"    {{ {row_vals} }},")
    lines.append("};")
    result = "\n".join(lines)

    os.makedirs("matrices", exist_ok=True)
    with open(f"matrices/{name}.h", "w") as f:
        f.write(result)

def save_cpp_array_1d(name, arr, dtype='float'):
    L = arr.shape[0]
    vals = ", ".join(cpp_float_repr(v) for v in arr)

    result = (
        f'static const {dtype} {name}[{L}] = {{\n'
        f"    {vals}\n"
        "};"
    )

    os.makedirs("matrices", exist_ok=True)
    with open(f"matrices/{name}.h", "w") as f:
        f.write(result)


def load_config(data_dir) -> Configuration:
    """Return typed access to the model and kernel sections in model.yaml."""
    document = load_model_document(data_dir)
    model = document.get('model', {})
    kernel = document.get('kernel', {})

    return Configuration(
        model=ModelConfig(
            sequence_length=int(model['sequence_length']),
            d_model=int(model['d_model']),
            num_heads=int(model['num_heads']),
            num_encoders=int(model['num_encoders']),
            ff_dim=int(model['ff_dim'])
        ),
        kernel=KernelConfig(
            store_fp16=bool(kernel['store_fp16']),
            acc_fp16=bool(kernel['acc_fp16']),
            encoder_layer_count=int(kernel.get('encoder_layer_count', 1))
        )
    )


def load_parameter_config(data_dir):
    """Return the full YAML document, including parameter extraction rules."""
    document = load_model_document(data_dir)
    return document
