import os
import torch
from torch import nn
import cv2
import numpy as np
from safetensors.torch import load_file, save_file

# архитектура
class VideoNet128(nn.Module):
    def __init__(self):
        super().__init__()
        self.enc = nn.Sequential(
            nn.Conv3d(3, 32, 3, padding=1), nn.ReLU(True), nn.MaxPool3d((1, 2, 2)),
            nn.Conv3d(32, 64, 3, padding=1), nn.ReLU(True), nn.MaxPool3d((1, 2, 2)),
            nn.Conv3d(64, 128, 3, padding=1), nn.ReLU(True), nn.MaxPool3d((1, 2, 2))
        )
        self.dec = nn.Sequential(
            nn.Upsample(scale_factor=(1, 2, 2)), nn.Conv3d(128, 64, 3, padding=1), nn.ReLU(True),
            nn.Upsample(scale_factor=(1, 2, 2)), nn.Conv3d(64, 32, 3, padding=1), nn.ReLU(True),
            nn.Upsample(scale_factor=(1, 2, 2)), nn.Conv3d(32, 16, 3, padding=1), nn.ReLU(True),
            nn.Conv3d(16, 3, 3, padding=1), nn.Tanh()
        )
    def forward(self, x): 
        return self.dec(self.enc(x))

# инициализация устройства
device = "cuda" if torch.cuda.is_available() else "cpu"
model = VideoNet128().to(device)

# конвертация и загрузка весов
pth_path = 'model.pth'
safetensors_path = 'model.safetensors'

if not os.path.exists(safetensors_path):
    if os.path.exists(pth_path):
        print("обнаружен старый формат весов! запускаю конвертацию в .safetensors...")
        try:
            model.load_state_dict(torch.load(pth_path, map_location='cpu'))
            save_file(model.state_dict(), safetensors_path)
            print("конвертация успешно завершена! создан файл model.safetensors.")
        except Exception as e:
            print(f"ошибка при конвертации: {e}")
    else:
        print(f"E: не найдены файлы весов ({pth_path} или {safetensors_path})!")
        exit(1)

# загрузка через safetensors
try:
    weights = load_file(safetensors_path, device=device)
    model.load_state_dict(weights)
    print("веса hueglot (.safetensors) успешно применены!")
except Exception as e:
    print(f"критическая ошибка загрузки .safetensors: {e}!")
    exit(1)

model.eval()

# обработка видео
def process_user_video(video_path, output_path="output.mp4"):
    # читаем видео и нарезаем первые 16 кадров в разрешении 128x128
    cap = cv2.VideoCapture(video_path)
    frames = []
    while len(frames) < 16:
        ret, frame = cap.read()
        if not ret: 
            break
        frame = cv2.resize(frame, (128, 128))
        frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
    cap.release()
    
    if len(frames) < 16: 
        print("E: видео слишком короткое, нужно минимум 16 кадров!")
        return

    # нормализация для hueglot (в диапазон от -1 до 1)
    v = np.array(frames, dtype=np.float32)
    v = torch.from_numpy(v).permute(3, 0, 1, 2).unsqueeze(0).to(device)
    v = (v / 127.5) - 1.0

    # прогоняем через нейросеть
    print("перемалываем видео через латентное пространство...")
    with torch.no_grad():
        output = model(v)
        vid = ((output.squeeze(0).clamp(-1, 1) + 1) / 2 * 255).cpu().numpy().astype(np.uint8).transpose(1, 2, 3, 0)

    # программный апскейл до 1024x1024 в 60 fps
    NEW_SIZE = (1024, 1024)
    out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), 60, NEW_SIZE)

    print("рендерим плавные 60 FPS через интерполяцию кадров...")
    for i in range(15):
        frame1 = cv2.resize(cv2.cvtColor(vid[i], cv2.COLOR_RGB2BGR), NEW_SIZE, interpolation=cv2.INTER_CUBIC)
        frame2 = cv2.resize(cv2.cvtColor(vid[i+1], cv2.COLOR_RGB2BGR), NEW_SIZE, interpolation=cv2.INTER_CUBIC)
        out.write(frame1)
        for alpha in [0.25, 0.5, 0.75]:
            out.write(cv2.addWeighted(frame1, 1 - alpha, frame2, alpha, 0))
            
    out.write(cv2.resize(cv2.cvtColor(vid[-1], cv2.COLOR_RGB2BGR), NEW_SIZE, interpolation=cv2.INTER_CUBIC))
    out.release()
    print(f"[+] Готово! Шизо-видео успешно обработано и сохранено в: {output_path}")

# ИСПОЛЬЗОВАНИЕ
# замени имя файла на свое входное видео
process_user_video("video.mp4")