{% extends "base.html" %}
{% block title %}IMPSY - Train{% endblock %}
{% block content %}
<h1 class="mb-2">Train</h1>
<p class="text-secondary mb-4">Train a model from a dataset. Training runs in the background and saves <code>.keras</code> and <code>.tflite</code> files to the models folder. A small model on a laptop CPU takes a few minutes.</p>

<div class="card mb-3">
    <div class="card-body">
        <h5 class="card-title">Train a Model</h5>
        {% if dataset_files %}
        <form method="POST" action="{{ url_for('train') }}" class="row g-2 align-items-end">
            <div class="col-12 col-md-auto">
                <label for="dataset" class="form-label">Dataset</label>
                <select class="form-select" id="dataset" name="dataset" required>
                    {% for file in dataset_files %}
                    <option value="{{ file.name }}" {% if file.name == selected_dataset %}selected{% endif %}>{{ file.name }}{% if file.dimension %} ({{ file.dimension }}d){% endif %}</option>
                    {% endfor %}
                </select>
            </div>
            <div class="col-auto">
                <label for="model_size" class="form-label">Model size</label>
                <select class="form-select" id="model_size" name="model_size">
                    {% for size in sizes %}
                    <option value="{{ size }}" {% if size == default_size %}selected{% endif %}>{{ size }}</option>
                    {% endfor %}
                </select>
            </div>
            <div class="col-auto">
                <label for="max_epochs" class="form-label">Max epochs</label>
                <input type="number" class="form-control" id="max_epochs" name="max_epochs" min="1" max="{{ max_epochs }}" value="100" style="width: 110px">
            </div>
            <div class="col-auto">
                <label for="patience" class="form-label">Patience</label>
                <input type="number" class="form-control" id="patience" name="patience" min="1" value="10" style="width: 100px">
            </div>
            <div class="col-auto">
                <button type="submit" class="btn btn-primary" {% if job.status in ['running', 'stopping'] %}disabled{% endif %}>Start Training</button>
            </div>
        </form>
        <div class="form-text mt-2">
            Training stops early when the validation loss hasn't improved for <em>patience</em> epochs.
            {% if config_dimension %}Your config uses dimension {{ config_dimension }}; the model's dimension comes from the dataset.{% endif %}
        </div>
        {% else %}
        <p class="text-secondary mb-0">No datasets yet. <a href="{{ url_for('datasets') }}">Generate a dataset</a> from your log files first.</p>
        {% endif %}
    </div>
</div>

<div class="card" id="job-card" {% if job.status == 'idle' %}hidden{% endif %}>
    <div class="card-body">
        <div class="d-flex justify-content-between align-items-start">
            <h5 class="card-title">Training Status</h5>
            <form method="POST" action="{{ url_for('train_stop') }}" id="stop-form" {% if job.status != 'running' %}hidden{% endif %}>
                <button type="submit" class="btn btn-outline-danger btn-sm">Stop and Save</button>
            </form>
        </div>
        <dl class="row mb-0">
            <dt class="col-sm-3">Status</dt><dd class="col-sm-9" id="job-status">{{ job.status }}</dd>
            <dt class="col-sm-3">Dataset</dt><dd class="col-sm-9"><code id="job-dataset">{{ job.dataset }}</code></dd>
            <dt class="col-sm-3">Model size</dt><dd class="col-sm-9" id="job-size">{{ job.model_size }}</dd>
            <dt class="col-sm-3">Epoch</dt><dd class="col-sm-9" id="job-epoch">{{ job.epoch }} / {{ job.max_epochs }}</dd>
            <dt class="col-sm-3">Loss (val)</dt><dd class="col-sm-9" id="job-loss">&ndash;</dd>
            <dt class="col-sm-3">Time</dt><dd class="col-sm-9" id="job-elapsed">{{ job.elapsed }} s</dd>
        </dl>
        <div class="alert alert-danger mt-3 mb-0" id="job-error" {% if not job.error %}hidden{% endif %}>{{ job.error }}</div>
        <div class="mt-3" id="job-files" {% if not job.model_files %}hidden{% endif %}>
            {% for name in job.model_files %}
            <a href="{{ url_for('download_model', filename=name) }}" class="btn btn-outline-secondary btn-sm me-1">Download {{ name }}</a>
            {% endfor %}
        </div>
    </div>
</div>
{% endblock %}

{% block scripts %}
<script>
(function () {
    const statusUrl = "{{ url_for('train_status') }}";
    const downloadUrl = "{{ url_for('download_model', filename='FILENAME') }}";
    const fmt = (v) => (v === null || v === undefined) ? "–" : v.toFixed(3);

    function render(job) {
        document.getElementById("job-card").hidden = job.status === "idle";
        document.getElementById("job-status").textContent = job.status;
        document.getElementById("job-dataset").textContent = job.dataset || "";
        document.getElementById("job-size").textContent = job.model_size || "";
        document.getElementById("job-epoch").textContent = job.epoch + " / " + job.max_epochs;
        const n = job.loss.length;
        document.getElementById("job-loss").textContent =
            n ? fmt(job.loss[n - 1]) + " (" + fmt(job.val_loss[n - 1]) + ")" : "–";
        document.getElementById("job-elapsed").textContent = job.elapsed + " s";
        document.getElementById("stop-form").hidden = job.status !== "running";
        const err = document.getElementById("job-error");
        err.hidden = !job.error;
        err.textContent = job.error || "";
        const files = document.getElementById("job-files");
        files.hidden = !job.model_files.length;
        files.replaceChildren(...job.model_files.map((name) => {
            const a = document.createElement("a");
            a.href = downloadUrl.replace("FILENAME", encodeURIComponent(name));
            a.className = "btn btn-outline-secondary btn-sm me-1";
            a.textContent = "Download " + name;
            return a;
        }));
    }

    let wasRunning = {{ 'true' if job.status in ['running', 'stopping'] else 'false' }};
    async function poll() {
        try {
            const job = await (await fetch(statusUrl)).json();
            render(job);
            const running = job.status === "running" || job.status === "stopping";
            if (wasRunning && !running) {
                // Re-enable the form now training has finished.
                location.reload();
                return;
            }
            wasRunning = running;
        } catch (e) { /* try again next time */ }
        setTimeout(poll, 2000);
    }
    poll();
})();
</script>
{% endblock %}