temp_backup.cpp

#include <windows.h>
#include <commdlg.h>
#include <shellapi.h>
#include <string>
#include <vector>
#include <thread>
#include <atomic>
#include <algorithm>
#include <random>
#include <ctime>
#include <shlwapi.h>
#include <shlobj.h>
#include <dwmapi.h>
#include <fstream>

#pragma comment(lib, "shlwapi.lib")
#pragma comment(lib, "comdlg32.lib")
#pragma comment(lib, "shell32.lib")
#pragma comment(lib, "dwmapi.lib")

struct Config {
    std::string inputFolder;
    std::string outputFolder;
    std::string rifePath;
    int numFrames = 1000;
    int epochs = 100;
    float learningRate = 0.01f;
    int batchSize = 32;
    bool cleanStatic = true;
    bool useRIFE = false;
};

struct Weights {
    float baseW1[4];
    float baseW2[4];
    float baseW3[4];
};

std::atomic<bool> g_training(false);
std::atomic<int> g_progress(0);
std::vector<std::string> g_log;

Config g_cfg;
Weights g_weights;

const COLORREF BG_COLOR = RGB(243, 243, 243);
const COLORREF CARD_COLOR = RGB(255, 255, 255);
const COLORREF ACCENT_COLOR = RGB(0, 120, 212);
const COLORREF ACCENT_LIGHT = RGB(153, 211, 253);
const COLORREF TEXT_COLOR = RGB(32, 32, 32);
const COLORREF TEXT_SECONDARY = RGB(136, 136, 136);
const COLORREF BORDER_COLOR = RGB(229, 229, 229);
const COLORREF INPUT_BG = RGB(251, 251, 251);
const COLORREF PROGRESS_BG = RGB(229, 229, 229);

bool fileExists(const std::string& path) {
    return GetFileAttributesA(path.c_str()) != INVALID_FILE_ATTRIBUTES;
}

void createFolder(const char* path) { CreateDirectoryA(path, nullptr); }

std::vector<std::string> getFiles(const std::string& folder, const std::vector<std::string>& exts) {
    std::vector<std::string> files;
    std::string searchPath = folder + "\\*.*";
    WIN32_FIND_DATAA fd;
    HANDLE hFind = FindFirstFileA(searchPath.c_str(), &fd);
    if (hFind != INVALID_HANDLE_VALUE) {
        do {
            if (!(fd.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY)) {
                std::string name = fd.cFileName;
                size_t dot = name.rfind('.');
                if (dot != std::string::npos) {
                    std::string ext = name.substr(dot + 1);
                    std::transform(ext.begin(), ext.end(), ext.begin(), ::tolower);
                    for (const auto& e : exts) {
                        if (ext == e) {
                            files.push_back(folder + "\\" + name);
                            break;
                        }
                    }
                }
            }
        } while (FindNextFileA(hFind, &fd));
        FindClose(hFind);
    }
    std::sort(files.begin(), files.end());
    return files;
}

void log(const std::string& msg) {
    g_log.push_back(msg);
    if (g_log.size() > 500) g_log.erase(g_log.begin());
}

void forward(const Weights& w, const float* in, float* out) {
    for (int i = 0; i < 4; i++) {
        float sum = 0;
        for (int j = 0; j < 12; j++) sum += in[j] * w.baseW1[j % 4];
        out[i] = 1.0f / (1.0f + exp(-sum));
    }
    for (int i = 0; i < 4; i++) {
        float sum = 0;
        for (int j = 0; j < 12; j++) sum += in[j] * w.baseW2[j % 4];
        out[4 + i] = 1.0f / (1.0f + exp(-sum));
    }
    for (int i = 0; i < 4; i++) {
        float sum = 0;
        for (int j = 0; j < 12; j++) sum += in[j] * w.baseW3[j % 4];
        out[8 + i] = 1.0f / (1.0f + exp(-sum));
    }
}

void initWeights(Weights& w) {
    for (int i = 0; i < 4; i++) {
        w.baseW1[i] = 0.5f + ((float)(rand() % 100) / 1000.0f);
        w.baseW2[i] = 0.5f + ((float)(rand() % 100) / 1000.0f);
        w.baseW3[i] = 0.5f + ((float)(rand() % 100) / 1000.0f);
    }
}

void cleanStaticImages(std::vector<std::string>& images) {
    if (images.size() < 3) return;
    std::vector<std::string> cleaned;
    for (size_t i = 1; i + 1 < images.size(); i++) cleaned.push_back(images[i]);
    images = cleaned;
}

void trainModel(HWND hwnd) {
    g_training = true;
    g_progress = 0;
    log("Starting training...");
    
    srand((unsigned int)time(nullptr));
    
    createFolder(g_cfg.outputFolder.c_str());
    createFolder((g_cfg.outputFolder + "\\frames").c_str());
    createFolder((g_cfg.outputFolder + "\\ground_truth").c_str());
    
    std::vector<std::string> imgExts = {"png", "jpg", "jpeg", "bmp", "tga"};
    std::vector<std::string> images = getFiles(g_cfg.inputFolder, imgExts);
    
    if (images.empty()) {
        log("ERROR: No images found!");
        g_training = false;
        return;
    }
    
    log("Found " + std::to_string(images.size()) + " images");
    
    if (g_cfg.cleanStatic && images.size() > 2) {
        cleanStaticImages(images);
        log("Using " + std::to_string(images.size()) + " frames after cleaning");
    }
    
    initWeights(g_weights);
    
    std::vector<std::vector<float>> features;
    std::vector<std::vector<float>> targets;
    
    std::random_device rd;
    std::mt19937 gen(rd());
    
    for (int i = 0; i < 1000; i++) {
        std::vector<float> feat(12);
        for (int j = 0; j < 12; j++) feat[j] = (float)(rd() % 100) / 100.0f;
        features.push_back(feat);
        
        std::vector<float> tgt(12, 0.25f);
        float maxF = *std::max_element(feat.begin(), feat.end());
        if (maxF > 0.5f) {
            int idx = (int)(std::max_element(feat.begin(), feat.end()) - feat.begin());
            tgt[idx] = 0.5f;
            tgt[(idx + 4) % 12] = 0.2f;
        }
        targets.push_back(tgt);
    }
    
    log("Training for " + std::to_string(g_cfg.epochs) + " epochs...");
    
    for (int epoch = 0; epoch < g_cfg.epochs && g_training; epoch++) {
        float totalLoss = 0;
        int count = 0;
        
        std::vector<int> idx(features.size());
        for (size_t i = 0; i < idx.size(); i++) idx[i] = (int)i;
        
        for (size_t i = idx.size() - 1; i > 0; i--) {
            int j = rd() % (i + 1);
            std::swap(idx[i], idx[j]);
        }
        
        for (size_t ii = 0; ii < idx.size() && g_training; ii++) {
            int i = idx[ii];
            float out[12] = { 0 };
            forward(g_weights, features[i].data(), out);
            
            float loss = 0;
            for (int j = 0; j < 12; j++) {
                float d = out[j] - targets[i][j];
                loss += d * d;
            }
            totalLoss += loss;
            count++;
            
            if (count % g_cfg.batchSize == 0) {
                for (int j = 0; j < 4; j++) {
                    g_weights.baseW1[j] -= g_cfg.learningRate * (out[j] - targets[i][j]) * 0.1f;
                    g_weights.baseW2[j] -= g_cfg.learningRate * (out[4 + j] - targets[i][4 + j]) * 0.1f;
                    g_weights.baseW3[j] -= g_cfg.learningRate * (out[8 + j] - targets[i][8 + j]) * 0.1f;
                    
                    if (g_weights.baseW1[j] < 0.01f) g_weights.baseW1[j] = 0.01f;
                    if (g_weights.baseW1[j] > 1.0f) g_weights.baseW1[j] = 1.0f;
                    if (g_weights.baseW2[j] < 0.01f) g_weights.baseW2[j] = 0.01f;
                    if (g_weights.baseW2[j] > 1.0f) g_weights.baseW2[j] = 1.0f;
                    if (g_weights.baseW3[j] < 0.01f) g_weights.baseW3[j] = 0.01f;
                    if (g_weights.baseW3[j] > 1.0f) g_weights.baseW3[j] = 1.0f;
                }
            }
        }
        
        g_progress = (int)((epoch + 1) * 100 / g_cfg.epochs);
        
        if (epoch % 10 == 0 || epoch == g_cfg.epochs - 1) {
            log("Epoch " + std::to_string(epoch + 1) + "/" + std::to_string(g_cfg.epochs) + 
                " | Loss: " + std::to_string(totalLoss / count));
        }
    }
    
    if (!g_training) {
        log("Training cancelled");
        g_training = false;
        return;
    }
    
    std::string outPath = g_cfg.outputFolder + "\\weights_trained.json";
    std::ofstream file(outPath);
    file << "{\n";
    file << "  \"baseW1\": [" << g_weights.baseW1[0] << ", " << g_weights.baseW1[1] << ", " 
         << g_weights.baseW1[2] << ", " << g_weights.baseW1[3] << "],\n";
    file << "  \"baseW2\": [" << g_weights.baseW2[0] << ", " << g_weights.baseW2[1] << ", " 
         << g_weights.baseW2[2] << ", " << g_weights.baseW2[3] << "],\n";
    file << "  \"baseW3\": [" << g_weights.baseW3[0] << ", " << g_weights.baseW3[1] << ", " 
         << g_weights.baseW3[2] << ", " << g_weights.baseW3[3] << "]\n";
    file << "}\n";
    file.close();
    
    log("SUCCESS! Saved to: " + outPath);
    g_training = false;
}

void drawRoundRect(HDC hdc, int x, int y, int w, int h, int r, COLORREF fill) {
    HRGN hrgn = CreateRoundRectRgn(x, y, x + w, y + h, r, r);
    HBRUSH hBrush = CreateSolidBrush(fill);
    FillRgn(hdc, hrgn, hBrush);
    DeleteObject(hBrush);
    DeleteObject(hrgn);
}

void drawRoundRectBorder(HDC hdc, int x, int y, int w, int h, int r, COLORREF fill, COLORREF border) {
    HRGN hrgn = CreateRoundRectRgn(x, y, x + w, y + h, r, r);
    HBRUSH hBrush = CreateSolidBrush(fill);
    FillRgn(hdc, hrgn, hBrush);
    DeleteObject(hBrush);
    
    HRGN borderRgn = CreateRoundRectRgn(x, y, x + w, y + h, r, r);
    HBRUSH hBorderBrush = CreateSolidBrush(border);
    FrameRgn(hdc, borderRgn, hBorderBrush, 1, 1);
    DeleteObject(hBorderBrush);
    DeleteObject(borderRgn);
    DeleteObject(hrgn);
}

void drawLabel(HDC hdc, int x, int y, const char* text, COLORREF color = TEXT_COLOR, int size = 14, bool bold = false) {
    SetBkMode(hdc, TRANSPARENT);
    SetTextColor(hdc, color);
    
    HFONT hFont = CreateFont(size, 0, 0, 0, bold ? FW_SEMIBOLD : FW_NORMAL, FALSE, FALSE, FALSE, DEFAULT_CHARSET, OUT_DEFAULT_PRECIS, CLIP_DEFAULT_PRECIS, CLEARTYPE_QUALITY, DEFAULT_PITCH, "Segoe UI Variable");
    HFONT hOld = (HFONT)SelectObject(hdc, hFont);
    
    RECT rc = { x, y, x + 400, y + 20 };
    DrawTextA(hdc, text, -1, &rc, DT_LEFT | DT_VCENTER | DT_SINGLELINE);
    
    SelectObject(hdc, hOld);
    DeleteObject(hFont);
}

void drawPrimaryButton(HDC hdc, int x, int y, int w, int h, const char* text, bool hover, bool pressed, bool disabled) {
    COLORREF bg = disabled ? RGB(204, 204, 204) : (pressed ? RGB(0, 95, 170) : (hover ? ACCENT_LIGHT : ACCENT_COLOR));
    drawRoundRect(hdc, x, y, w, h, 6, bg);
    
    SetBkMode(hdc, TRANSPARENT);
    SetTextColor(hdc, disabled ? RGB(150, 150, 150) : RGB(255, 255, 255));
    
    HFONT hFont = CreateFont(14, 0, 0, 0, FW_MEDIUM, FALSE, FALSE, FALSE, DEFAULT_CHARSET, OUT_DEFAULT_PRECIS, CLIP_DEFAULT_PRECIS, CLEARTYPE_QUALITY, DEFAULT_PITCH, "Segoe UI Variable");
    HFONT hOld = (HFONT)SelectObject(hdc, hFont);
    
    RECT rc = { x, y, x + w, y + h };
    DrawTextA(hdc, text, -1, &rc, DT_CENTER | DT_VCENTER | DT_SINGLELINE);
    
    SelectObject(hdc, hOld);
    DeleteObject(hFont);
}

void drawSecondaryButton(HDC hdc, int x, int y, int w, int h, const char* text, bool hover, bool pressed, bool disabled) {
    COLORREF bg = pressed ? RGB(230, 230, 230) : (hover ? RGB(246, 246, 246) : CARD_COLOR);
    COLORREF border = disabled ? RGB(204, 204, 204) : (hover ? ACCENT_COLOR : BORDER_COLOR);
    drawRoundRectBorder(hdc, x, y, w, h, 6, bg, border);
    
    SetBkMode(hdc, TRANSPARENT);
    SetTextColor(hdc, disabled ? RGB(150, 150, 150) : TEXT_COLOR);
    
    HFONT hFont = CreateFont(14, 0, 0, 0, FW_MEDIUM, FALSE, FALSE, FALSE, DEFAULT_CHARSET, OUT_DEFAULT_PRECIS, CLIP_DEFAULT_PRECIS, CLEARTYPE_QUALITY, DEFAULT_PITCH, "Segoe UI Variable");
    HFONT hOld = (HFONT)SelectObject(hdc, hFont);
    
    RECT rc = { x, y, x + w, y + h };
    DrawTextA(hdc, text, -1, &rc, DT_CENTER | DT_VCENTER | DT_SINGLELINE);
    
    SelectObject(hdc, hOld);
    DeleteObject(hFont);
}

void drawProgressBar(HDC hdc, int x, int y, int w, int h, int percent) {
    drawRoundRect(hdc, x, y, w, h, 4, PROGRESS_BG);
    
    if (percent > 0) {
        int barW = (int)((w - 4) * percent / 100.0f);
        if (barW > 0) {
            drawRoundRect(hdc, x + 2, y + 2, barW, h - 4, 3, ACCENT_COLOR);
        }
    }
}

void drawCheckBox(HDC hdc, int x, int y, const char* text, bool checked, bool hover) {
    COLORREF boxBg = checked ? ACCENT_COLOR : (hover ? RGB(240, 240, 240) : CARD_COLOR);
    drawRoundRectBorder(hdc, x, y + 2, 18, 18, 3, boxBg, checked ? ACCENT_COLOR : BORDER_COLOR);
    
    if (checked) {
        SetBkMode(hdc, TRANSPARENT);
        SetTextColor(hdc, RGB(255, 255, 255));
        HFONT hFont = CreateFont(10, 0, 0, 0, FW_BOLD, FALSE, FALSE, FALSE, DEFAULT_CHARSET, OUT_DEFAULT_PRECIS, CLIP_DEFAULT_PRECIS, CLEARTYPE_QUALITY, DEFAULT_PITCH, "Segoe UI");
        HFONT hOldF = (HFONT)SelectObject(hdc, hFont);
        RECT rc = { x, y + 2, x + 18, y + 20 };
        DrawTextA(hdc, "OK", -1, &rc, DT_CENTER | DT_VCENTER);
        SelectObject(hdc, hOldF);
        DeleteObject(hFont);
    }
    
    drawLabel(hdc, x + 26, y + 1, text);
}

HWND hInputEdit, hOutputEdit, hEpochsEdit, hLREdit;
int g_hoverBtn = 0;
bool g_pressedBtn = false;

void invalidateAll(HWND hwnd) {
    InvalidateRect(hwnd, NULL, FALSE);
}

void redrawWindow(HWND hwnd, HDC hdc) {
    RECT rc;
    GetClientRect(hwnd, &rc);
    
    HBRUSH hBg = CreateSolidBrush(BG_COLOR);
    FillRect(hdc, &rc, hBg);
    DeleteObject(hBg);
    
    HFONT hTitleFont = CreateFont(28, 0, 0, 0, FW_SEMIBOLD, FALSE, FALSE, FALSE, DEFAULT_CHARSET, OUT_DEFAULT_PRECIS, CLIP_DEFAULT_PRECIS, CLEARTYPE_QUALITY, DEFAULT_PITCH, "Segoe UI Variable");
    HFONT hOld = (HFONT)SelectObject(hdc, hTitleFont);
    SetTextColor(hdc, TEXT_COLOR);
    SetBkMode(hdc, TRANSPARENT);
    RECT titleRc = { 24, 20, 500, 60 };
    DrawTextA(hdc, "True Motion Fidelity", -1, &titleRc, DT_LEFT | DT_TOP | DT_SINGLELINE);
    SelectObject(hdc, hOld);
    DeleteObject(hTitleFont);
    
    drawLabel(hdc, 24, 55, "Training Suite", TEXT_SECONDARY, 14);
    
    int cardY = 80;
    drawRoundRect(hdc, 16, cardY, 508, 140, 8, CARD_COLOR);
    
    drawLabel(hdc, 32, cardY + 16, "Input Folder", TEXT_SECONDARY, 13);
    
    char inBuf[260] = { 0 };
    GetWindowText(hInputEdit, inBuf, 259);
    drawRoundRectBorder(hdc, 32, cardY + 34, 360, 32, 4, INPUT_BG, g_hoverBtn == 1 ? ACCENT_COLOR : BORDER_COLOR);
    drawLabel(hdc, 40, cardY + 40, inBuf[0] ? inBuf : "Select folder containing images or videos...", 
              inBuf[0] ? TEXT_COLOR : TEXT_SECONDARY, 13);
    drawSecondaryButton(hdc, 400, cardY + 34, 100, 32, "Browse", g_hoverBtn == 1, g_pressedBtn && g_hoverBtn == 1, FALSE);
    
    drawLabel(hdc, 32, cardY + 76, "Output Folder", TEXT_SECONDARY, 13);
    
    char outBuf[260] = { 0 };
    GetWindowText(hOutputEdit, outBuf, 259);
    drawRoundRectBorder(hdc, 32, cardY + 94, 360, 32, 4, INPUT_BG, BORDER_COLOR);
    drawLabel(hdc, 40, cardY + 100, outBuf[0] ? outBuf : "Select output folder...", 
              outBuf[0] ? TEXT_COLOR : TEXT_SECONDARY, 13);
    drawSecondaryButton(hdc, 400, cardY + 94, 100, 32, "Browse", g_hoverBtn == 2, g_pressedBtn && g_hoverBtn == 2, FALSE);
    
    cardY += 148;
    drawRoundRect(hdc, 16, cardY, 508, 110, 8, CARD_COLOR);
    
    drawLabel(hdc, 32, cardY + 16, "Training Settings", TEXT_SECONDARY, 13);
    
    drawLabel(hdc, 32, cardY + 42, "Epochs", TEXT_COLOR, 13);
    char epBuf[32] = { 0 };
    GetWindowText(hEpochsEdit, epBuf, 31);
    drawRoundRectBorder(hdc, 90, cardY + 40, 80, 28, 4, INPUT_BG, BORDER_COLOR);
    drawLabel(hdc, 96, cardY + 45, epBuf[0] ? epBuf : "100", TEXT_COLOR, 13);
    
    drawLabel(hdc, 200, cardY + 42, "Learning Rate", TEXT_COLOR, 13);
    char lrBuf[32] = { 0 };
    GetWindowText(hLREdit, lrBuf, 31);
    drawRoundRectBorder(hdc, 310, cardY + 40, 80, 28, 4, INPUT_BG, BORDER_COLOR);
    drawLabel(hdc, 316, cardY + 45, lrBuf[0] ? lrBuf : "0.01", TEXT_COLOR, 13);
    
    bool chkClean = (SendMessage(GetDlgItem(hwnd, 400), BM_GETCHECK, 0, 0) == BST_CHECKED);
    drawCheckBox(hdc, 32, cardY + 78, "Clean static images", chkClean, false);
    
    if (!g_training) {
        drawPrimaryButton(hdc, 400, cardY + 68, 108, 36, "Start Training", g_hoverBtn == 3, g_pressedBtn && g_hoverBtn == 3, FALSE);
    } else {
        drawSecondaryButton(hdc, 400, cardY + 68, 108, 36, "Stop", g_hoverBtn == 4, g_pressedBtn && g_hoverBtn == 4, FALSE);
    }
    
    int progY = cardY + 125;
    drawLabel(hdc, 32, progY, "Progress", TEXT_COLOR, 14, true);
    drawProgressBar(hdc, 32, progY + 24, 436, 8, g_progress.load());
    char pct[16];
    sprintf(pct, "%d%%", g_progress.load());
    drawLabel(hdc, 476, progY + 20, pct, TEXT_SECONDARY, 13);
    
    int logY = progY + 60;
    drawRoundRect(hdc, 16, logY, 508, 130, 8, CARD_COLOR);
    
    drawLabel(hdc, 32, logY + 12, "Output Log", TEXT_SECONDARY, 13);
    
    std::string allLog;
    for (const auto& s : g_log) {
        if (allLog.size() > 2000) break;
        allLog += s + "\n";
    }
    
    RECT logBox = { 28, logY + 32, 496, logY + 112 };
    drawRoundRect(hdc, logBox.left, logBox.top, logBox.right - logBox.left, logBox.bottom - logBox.top, 4, INPUT_BG);
    
    if (!allLog.empty()) {
        SetBkMode(hdc, TRANSPARENT);
        SetTextColor(hdc, TEXT_COLOR);
        HFONT hLogFont = CreateFont(12, 0, 0, 0, FW_NORMAL, FALSE, FALSE, FALSE, DEFAULT_CHARSET, OUT_DEFAULT_PRECIS, CLIP_DEFAULT_PRECIS, CLEARTYPE_QUALITY, DEFAULT_PITCH, "Cascadia Mono");
        HFONT hOldL = (HFONT)SelectObject(hdc, hLogFont);
        RECT textRc = { logBox.left + 8, logBox.top + 4, logBox.right - 8, logBox.bottom };
        DrawTextA(hdc, allLog.c_str(), -1, &textRc, DT_LEFT | DT_TOP | DT_WORDBREAK | DT_NOCLIP);
        SelectObject(hdc, hOldL);
        DeleteObject(hLogFont);
    }
}

INT_PTR CALLBACK DialogProc(HWND hwnd, UINT msg, WPARAM wParam, LPARAM lParam) {
    switch (msg) {
        case WM_INITDIALOG: {
            SetWindowPos(hwnd, nullptr, 0, 0, 560, 620, SWP_NOMOVE | SWP_NOZORDER | SWP_FRAMECHANGED);
            
            hInputEdit = CreateWindowExA(0, "EDIT", "", WS_VISIBLE | WS_CHILD | WS_BORDER | ES_AUTOHSCROLL | ES_READONLY, 
                32, 114, 360, 32, hwnd, (HMENU)100, NULL, NULL);
            hOutputEdit = CreateWindowExA(0, "EDIT", "", WS_VISIBLE | WS_CHILD | WS_BORDER | ES_AUTOHSCROLL | ES_READONLY, 
                32, 174, 360, 32, hwnd, (HMENU)101, NULL, NULL);
            hEpochsEdit = CreateWindowExA(0, "EDIT", "", WS_VISIBLE | WS_CHILD | WS_BORDER | ES_AUTOHSCROLL | ES_NUMBER, 
                90, 164, 80, 28, hwnd, (HMENU)102, NULL, NULL);
            hLREdit = CreateWindowExA(0, "EDIT", "", WS_VISIBLE | WS_CHILD | WS_BORDER | ES_AUTOHSCROLL, 
                310, 164, 80, 28, hwnd, (HMENU)103, NULL, NULL);
            
            HWND hCheck = CreateWindowExA(0, "BUTTON", "Clean static images", WS_VISIBLE | WS_CHILD | BS_AUTOCHECKBOX, 
                32, 200, 150, 24, hwnd, (HMENU)400, NULL, NULL);
            SendMessage(hCheck, BM_SETCHECK, BST_CHECKED, 0);
            
            SetWindowText(hInputEdit, "dataset");
            SetWindowText(hOutputEdit, "output");
            SetWindowText(hEpochsEdit, "100");
            SetWindowText(hLREdit, "0.01");
            
            SetTimer(hwnd, 1, 100, NULL);
            return TRUE;
        }
        
        case WM_ERASEBKGND:
            return TRUE;
            
        case WM_PAINT: {
            PAINTSTRUCT ps;
            HDC hdc = BeginPaint(hwnd, &ps);
            redrawWindow(hwnd, hdc);
            EndPaint(hwnd, &ps);
            return TRUE;
        }
        
        case WM_MOUSEMOVE: {
            int x = LOWORD(lParam);
            int y = HIWORD(lParam);
            
            int oldHover = g_hoverBtn;
            g_hoverBtn = 0;
            
            if (x >= 400 && x <= 500 && y >= 114 && y <= 146) g_hoverBtn = 1;
            else if (x >= 400 && x <= 500 && y >= 174 && y <= 206) g_hoverBtn = 2;
            else if (!g_training && x >= 400 && x <= 508 && y >= 248 && y <= 284) g_hoverBtn = 3;
            else if (g_training && x >= 400 && x <= 508 && y >= 248 && y <= 284) g_hoverBtn = 4;
            
            if (oldHover != g_hoverBtn) invalidateAll(hwnd);
            
            if (g_pressedBtn && g_hoverBtn == 0) {
                g_pressedBtn = false;
                invalidateAll(hwnd);
            }
            break;
        }
        
        case WM_LBUTTONDOWN: {
            int x = LOWORD(lParam);
            int y = HIWORD(lParam);
            
            if (x >= 400 && x <= 500 && y >= 114 && y <= 146) {
                g_pressedBtn = true;
                g_hoverBtn = 1;
            } else if (x >= 400 && x <= 500 && y >= 174 && y <= 206) {
                g_pressedBtn = true;
                g_hoverBtn = 2;
            } else if (!g_training && x >= 400 && x <= 508 && y >= 248 && y <= 284) {
                g_pressedBtn = true;
                g_hoverBtn = 3;
            } else if (g_training && x >= 400 && x <= 508 && y >= 248 && y <= 284) {
                g_pressedBtn = true;
                g_hoverBtn = 4;
            } else {
                SetFocus(hwnd);
            }
            invalidateAll(hwnd);
            break;
        }
        
        case WM_LBUTTONUP: {
            int x = LOWORD(lParam);
            int y = HIWORD(lParam);
            
            if (g_pressedBtn && g_hoverBtn == 1) {
                char path[MAX_PATH] = { 0 };
                BROWSEINFOA bi = { 0 };
                bi.hwndOwner = hwnd;
                bi.lpszTitle = "Select Input Folder";
                bi.ulFlags = BIF_RETURNONLYFSDIRS;
                LPITEMIDLIST pidl = SHBrowseForFolderA(&bi);
                if (pidl && SHGetPathFromIDListA(pidl, path)) {
                    SetWindowText(hInputEdit, path);
                }
            } else if (g_pressedBtn && g_hoverBtn == 2) {
                char path[MAX_PATH] = { 0 };
                BROWSEINFOA bi = { 0 };
                bi.hwndOwner = hwnd;
                bi.lpszTitle = "Select Output Folder";
                bi.ulFlags = BIF_RETURNONLYFSDIRS;
                LPITEMIDLIST pidl = SHBrowseForFolderA(&bi);
                if (pidl && SHGetPathFromIDListA(pidl, path)) {
                    SetWindowText(hOutputEdit, path);
                }
            } else if (g_pressedBtn && g_hoverBtn == 3 && !g_training) {
                char buf[256];
                GetWindowText(hInputEdit, buf, 255);
                g_cfg.inputFolder = buf;
                GetWindowText(hOutputEdit, buf, 255);
                g_cfg.outputFolder = buf;
                GetWindowText(hEpochsEdit, buf, 255);
                g_cfg.epochs = atoi(buf);
                GetWindowText(hLREdit, buf, 255);
                g_cfg.learningRate = (float)atof(buf);
                g_cfg.cleanStatic = (SendMessage(GetDlgItem(hwnd, 400), BM_GETCHECK, 0, 0) == BST_CHECKED);
                
                g_log.clear();
                log("=== True Motion Fidelity Training ===");
                log("Input: " + g_cfg.inputFolder);
                log("Output: " + g_cfg.outputFolder);
                log("Epochs: " + std::to_string(g_cfg.epochs) + " | LR: " + std::to_string(g_cfg.learningRate));
                
                std::thread([hwnd]() {
                    trainModel(hwnd);
                }).detach();
            } else if (g_pressedBtn && g_hoverBtn == 4 && g_training) {
                g_training = false;
                log("Stopping...");
            }
            
            g_pressedBtn = false;
            g_hoverBtn = 0;
            invalidateAll(hwnd);
            break;
        }
            
        case WM_TIMER:
            invalidateAll(hwnd);
            break;
            
        case WM_CLOSE:
            g_training = false;
            EndDialog(hwnd, 0);
            return TRUE;
    }
    return FALSE;
}

int WINAPI WinMain(HINSTANCE hInst, HINSTANCE, LPSTR, int) {
    HWND hwnd = CreateDialog(hInst, MAKEINTRESOURCE(1), NULL, DialogProc);
    if (!hwnd) return 1;
    
    ShowWindow(hwnd, SW_SHOW);
    UpdateWindow(hwnd);
    
    MSG msg;
    while (GetMessage(&msg, NULL, 0, 0)) {
        TranslateMessage(&msg);
        DispatchMessage(&msg);
    }
    return 0;
}