from jinja2 import BaseLoader, Environment

_ENV = Environment(loader=BaseLoader(), trim_blocks=True, lstrip_blocks=True)


def _render(template, **context):
    rendered = _ENV.from_string(template).render(**context)
    return rendered.rstrip() + "\n"


CONFIG_TEMPLATE = '''CONFIG = {
    "model_name": "{{ name }}",
    "input_dim": {{ input_dim }},
    "hidden": {{ hidden }},
    "output_dim": {{ output_dim }},
    "lr": {{ lr }},
    "epochs": {{ epochs }},
    "batch_size": {{ batch_size }},
    "device": "{{ device }}",
}
'''

MODEL_TEMPLATE = '''import torch.nn as nn


class {{ name }}(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
{{ layers_code }}
        )

    def forward(self, x):
        return self.net(x)
'''

DATA_TEMPLATE = '''import torch
from torch.utils.data import Dataset, DataLoader, random_split
from .config import CONFIG


class {{ name }}Dataset(Dataset):
    def __init__(self, X, y):
        self.X = torch.tensor(X, dtype=torch.float32)
        self.y = torch.tensor(y, dtype=torch.long)

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]


def get_loaders(X, y, val_split=0.2):
    dataset = {{ name }}Dataset(X, y)
    val_size = int(len(dataset) * val_split)
    train_size = len(dataset) - val_size
    train_ds, val_ds = random_split(dataset, [train_size, val_size])

    train_loader = DataLoader(train_ds, batch_size=CONFIG["batch_size"], shuffle=True)
    val_loader = DataLoader(val_ds, batch_size=CONFIG["batch_size"], shuffle=False)
    return train_loader, val_loader
'''

TRAIN_TEMPLATE = '''import torch
import torch.nn as nn
{% if has_amp %}from torch.amp import GradScaler, autocast
{% endif %}{% if prod %}import mlflow
{% else %}from torch.utils.tensorboard import SummaryWriter
{% endif %}{% if needs_sklearn %}from sklearn.metrics import {{ sklearn_imports|join(', ') }}
{% endif %}from .models import {{ name }}
from .config import CONFIG

device = torch.device("{{ device }}")
model = torch.compile({{ name }}().to(device))
optimizer = torch.optim.Adam(model.parameters(), lr=CONFIG["lr"])
criterion = nn.CrossEntropyLoss()
{% if has_amp %}scaler = GradScaler("{{ device }}")
{% endif %}{% if not prod %}writer = SummaryWriter("runs/" + CONFIG["model_name"])
{% endif %}


def train(train_loader, val_loader=None):
    if val_loader is None:
        val_loader = train_loader
{% if prod %}
    mlflow.set_experiment(CONFIG["model_name"])
    with mlflow.start_run():
        mlflow.log_params(CONFIG)
        for epoch in range(CONFIG["epochs"]):
            model.train()
            train_loss = 0.0
            for data, target in train_loader:
                data, target = data.to(device), target.to(device)
                optimizer.zero_grad()
{% if has_amp %}
                with autocast("{{ device }}"):
                    loss = criterion(model(data), target)
                scaler.scale(loss).backward()
                scaler.step(optimizer)
                scaler.update()
{% else %}
                loss = criterion(model(data), target)
                loss.backward()
                optimizer.step()
{% endif %}
                train_loss += loss.item()

            train_loss /= len(train_loader)
            mlflow.log_metric("train_loss", train_loss, step=epoch)

{% if has_val %}
            model.eval()
            val_loss = 0.0
            correct = 0
            total = 0
{% if needs_sklearn %}
            all_preds = []
            all_targets = []
{% endif %}
            with torch.no_grad():
                for data, target in val_loader:
                    data, target = data.to(device), target.to(device)
                    output = model(data)
                    val_loss += criterion(output, target).item()
                    _, predicted = output.max(1)
                    total += target.size(0)
                    correct += predicted.eq(target).sum().item()
{% if needs_sklearn %}
                    all_preds.extend(predicted.cpu().tolist())
                    all_targets.extend(target.cpu().tolist())
{% endif %}

            val_loss /= len(val_loader)
            accuracy = 100.0 * correct / total
            mlflow.log_metric("val_loss", val_loss, step=epoch)
            mlflow.log_metric("val_accuracy", accuracy, step=epoch)
{% if "f1" in metrics %}
            f1 = f1_score(all_targets, all_preds, average="weighted")
            mlflow.log_metric("val_f1", f1, step=epoch)
{% endif %}{% if "precision" in metrics %}
            prec = precision_score(all_targets, all_preds, average="weighted")
            mlflow.log_metric("val_precision", prec, step=epoch)
{% endif %}{% if "recall" in metrics %}
            rec = recall_score(all_targets, all_preds, average="weighted")
            mlflow.log_metric("val_recall", rec, step=epoch)
{% endif %}

            print(
                f"Epoch {epoch+1}/{CONFIG['epochs']} | Train Loss: {train_loss:.4f} | "
                f"Val Loss: {val_loss:.4f} | Acc: {accuracy:.2f}%"
            )
{% else %}
            print(
                f"Epoch {epoch+1}/{CONFIG['epochs']} | Train Loss: {train_loss:.4f}"
            )
{% endif %}

        torch.save(model.state_dict(), "model.pt")
        mlflow.log_artifact("model.pt")
{% else %}
    for epoch in range(CONFIG["epochs"]):
        model.train()
        train_loss = 0.0
        for data, target in train_loader:
            data, target = data.to(device), target.to(device)
            optimizer.zero_grad()
{% if has_amp %}
            with autocast("{{ device }}"):
                loss = criterion(model(data), target)
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
{% else %}
            loss = criterion(model(data), target)
            loss.backward()
            optimizer.step()
{% endif %}
            train_loss += loss.item()

        train_loss /= len(train_loader)
        writer.add_scalar("Loss/train", train_loss, epoch)

{% if has_val %}
        model.eval()
        val_loss = 0.0
        correct = 0
        total = 0
{% if needs_sklearn %}
        all_preds = []
        all_targets = []
{% endif %}
        with torch.no_grad():
            for data, target in val_loader:
                data, target = data.to(device), target.to(device)
                output = model(data)
                val_loss += criterion(output, target).item()
                _, predicted = output.max(1)
                total += target.size(0)
                correct += predicted.eq(target).sum().item()
{% if needs_sklearn %}
                all_preds.extend(predicted.cpu().tolist())
                all_targets.extend(target.cpu().tolist())
{% endif %}

        val_loss /= len(val_loader)
        accuracy = 100.0 * correct / total
        writer.add_scalar("Loss/val", val_loss, epoch)
        writer.add_scalar("Accuracy/val", accuracy, epoch)
{% if "f1" in metrics %}
        f1 = f1_score(all_targets, all_preds, average="weighted")
        writer.add_scalar("F1/val", f1, epoch)
{% endif %}{% if "precision" in metrics %}
        prec = precision_score(all_targets, all_preds, average="weighted")
        writer.add_scalar("Precision/val", prec, epoch)
{% endif %}{% if "recall" in metrics %}
        rec = recall_score(all_targets, all_preds, average="weighted")
        writer.add_scalar("Recall/val", rec, epoch)
{% endif %}

        print(
            f"Epoch {epoch+1}/{CONFIG['epochs']} | Train Loss: {train_loss:.4f} | "
            f"Val Loss: {val_loss:.4f} | Acc: {accuracy:.2f}%"
        )
{% else %}
        print(
            f"Epoch {epoch+1}/{CONFIG['epochs']} | Train Loss: {train_loss:.4f}"
        )
{% endif %}

    writer.close()
    torch.save(model.state_dict(), "model.pt")
{% endif %}
'''

PREDICT_TEMPLATE = '''import torch
from .models import {{ name }}
from .config import CONFIG

device = torch.device(CONFIG["device"])
model = {{ name }}().to(device)
model.load_state_dict(torch.load("model.pt", map_location=device))
model.eval()


def predict(x):
    with torch.no_grad():
        return model(x.to(device))
'''

INIT_TEMPLATE = '''from .models import {{ name }}
from .config import CONFIG
from .data import {{ name }}Dataset, get_loaders
from .train import train
from .predict import predict

__all__ = ["{{ name }}", "CONFIG", "{{ name }}Dataset", "get_loaders", "train", "predict"]
'''

MAIN_TEMPLATE = '''import torch
from src.config import CONFIG
from src.data import get_loaders
from src.train import train
from src.predict import predict


def main():
    X = torch.randn(512, CONFIG["input_dim"])
    y = torch.randint(0, CONFIG["output_dim"], (512,))
    train_loader, val_loader = get_loaders(X, y)
    train(train_loader, val_loader)

    sample = torch.randn(4, CONFIG["input_dim"])
    outputs = predict(sample)
    print("Sample outputs:", outputs)


if __name__ == "__main__":
    main()
'''


def render_config(name, input_dim, hidden, output_dim, lr, epochs, batch_size, device):
    return _render(
        CONFIG_TEMPLATE,
        name=name,
        input_dim=input_dim,
        hidden=hidden,
        output_dim=output_dim,
        lr=lr,
        epochs=epochs,
        batch_size=batch_size,
        device=device,
    )


def render_model(name, layers_code):
    return _render(MODEL_TEMPLATE, name=name, layers_code=layers_code)


def render_data(name):
    return _render(DATA_TEMPLATE, name=name)


def render_train(name, device, prod, metrics, has_amp):
    metrics = metrics or []
    has_val = len(metrics) > 0
    sklearn_metrics = set(metrics) & {"f1", "precision", "recall"}
    sklearn_imports = []
    if "f1" in metrics:
        sklearn_imports.append("f1_score")
    if "precision" in metrics:
        sklearn_imports.append("precision_score")
    if "recall" in metrics:
        sklearn_imports.append("recall_score")
    return _render(
        TRAIN_TEMPLATE,
        name=name,
        device=device,
        prod=prod,
        metrics=metrics,
        has_amp=has_amp,
        has_val=has_val,
        needs_sklearn=bool(sklearn_metrics),
        sklearn_imports=sklearn_imports,
    )


def render_predict(name):
    return _render(PREDICT_TEMPLATE, name=name)


def render_init(name):
    return _render(INIT_TEMPLATE, name=name)


def render_main():
    return _render(MAIN_TEMPLATE)


def render_requirements(prod, metrics):
    deps = ["torch>=2.1"]
    metrics = metrics or []
    if prod:
        deps.append("mlflow>=3.7")
    else:
        deps.append("tensorboard>=2.18")
    if set(metrics) & {"f1", "precision", "recall"}:
        deps.append("scikit-learn>=1.4")
    return "\n".join(deps) + "\n"
