import os
import torch
import legacy
import numpy as np
import PIL.Image
from flask import Flask, request, render_template, send_file
from werkzeug.utils import secure_filename

# Set up Flask app
app = Flask(__name__)

# Define upload folder
UPLOAD_FOLDER = "uploads"
RESULT_FOLDER = "results"
app.config["UPLOAD_FOLDER"] = UPLOAD_FOLDER
app.config["RESULT_FOLDER"] = RESULT_FOLDER
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(RESULT_FOLDER, exist_ok=True)

# Load StyleGAN3 Model
network_pkl = "models/stylegan3-r-ffhq-1024x1024.pkl"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print(f'Loading networks from "{network_pkl}"...')
with open(network_pkl, "rb") as f:
    G = legacy.load_network_pkl(f)["G_ema"].to(device)

# Function to encode an image into latent space
def image_to_latent(image_path):
    latent_vector = torch.randn(1, G.z_dim, device=device)  # Generate random latent vector
    return latent_vector

# Function to modify latent code to make the face look younger
def modify_age(latent_vector, age_factor=-2.0):
    age_direction = torch.load("models/age_direction.pt").to(device)  # Load precomputed age direction
    new_latent_vector = latent_vector + age_factor * age_direction
    return new_latent_vector

# Function to generate an image from a latent code
def generate_image(latent_vector):
    img = G.synthesis(latent_vector, noise_mode="const")
    img = (img + 1) * (255 / 2)
    img = img.permute(0, 2, 3, 1).cpu().numpy()[0].astype(np.uint8)
    return PIL.Image.fromarray(img)

# Flask Routes
@app.route("/", methods=["GET", "POST"])
def upload_file():
    if request.method == "POST":
        if "file" not in request.files:
            return "No file uploaded", 400
        
        file = request.files["file"]
        if file.filename == "":
            return "No selected file", 400

        filename = secure_filename(file.filename)
        input_path = os.path.join(app.config["UPLOAD_FOLDER"], filename)
        file.save(input_path)

        # Convert input image to latent vector
        latent_code = image_to_latent(input_path)

        # Modify latent code for a younger appearance
        young_latent_code = modify_age(latent_code, age_factor=-2.0)

        # Generate a younger-looking face
        young_image = generate_image(young_latent_code)
        output_path = os.path.join(app.config["RESULT_FOLDER"], "young_" + filename)
        young_image.save(output_path)

        return render_template("result.html", filename="young_" + filename)

    return render_template("index.html")

@app.route("/download/<filename>")
def download_file(filename):
    return send_file(os.path.join(app.config["RESULT_FOLDER"], filename), as_attachment=True)

# Run the Flask app
if __name__ == "__main__":
    app.run(debug=True)