import spaces  # oltava ensimmäisenä!

import torch
from diffusers import StableDiffusionXLPipeline
import gradio as gr
import random
import re

model_id = "kkntr/sdxl-kkntr"

pipe = StableDiffusionXLPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    use_safetensors=True
).to("cuda")

pipe.load_lora_weights(
    {
        "kkntr/lora1": 0.2,
        "kkntr/lora2": 0.3,
        "kkntr/lora12": 0.4,
        "kkntr/lora10": 0.1,
    }
)
pipe.fuse_lora()

# Scratchpad vaihtoehdot
scratchpad = {
    "animal": [
        "kangaroo", "kangaroo", "otter", "otter", "shark", "shark", "shark", "dolphin", "dolphin", "avian", "tiger", "lion", "wolf", "fox", "polar bear", "polar bear", "deer", "deer", "reindeer", "reindeer", "hyena", "bull", "bull", 
"rabbit", "panther", "moose", "cheetah", "badger", "striped hyena", "rat", "leopard", "thylacine", "bighorn sheep", 
"zebra", "horse", "horse", "donkey", "donkey", "unicorn", "white unicorn", "alligator", "western dragon", "lizard", "crocodile", "theropod",
"kangaroo", "otter", "tiger", "lion", "wolf", "fox", "polar bear", "deer", "reindeer", "hyena", "bull", 
"rabbit", "panther", "moose", "cheetah", "badger", "striped hyena", "rat", "leopard", "thylacine", "bighorn sheep", 
"zebra", "horse", "donkey", "unicorn", "white unicorn", "alligator", "western dragon", "lizard", "crocodile", "theropod",
"kangaroo", "otter", "tiger", "lion", "wolf", "fox", "polar bear", "deer", "reindeer", "hyena", "bull", 
"rabbit", "panther", "moose", "cheetah", "badger", "striped hyena", "rat", "leopard", "thylacine", "bighorn sheep", 
"zebra", "horse", "donkey", "unicorn", "white unicorn", "alligator", "western dragon", "lizard", "crocodile", "theropod",
"black russian terrier", "pit bull", "schnauzer", "dobermann", "mastiff", "molosser", "pinscher", "rottweiler", "tamaskan dog", 
"wolfdog", "alaskan husky", "alaskan malamute", "canadian eskimo dog", "german shepherd", "labrador husky", "siberian husky", 
"labrador", "golden retriever", "dalmatian", "border collie", "collie", "herding dog", "pastoral dog", "sheepdog", 
"carolina dog", "rhodesian ridgeback", "nordic sled dog", "spitz", "husky", "akita", "mountain dog", "primitive dog", 
"siberian retriever", "hunting dog", "australian cattle dog", "crash bandicoot", "montgomery gator, fnaf", "bear",
"parrot", "bison", "buffalo", "cattle", "muskox", "yak", "alpaca", "goat", "sheep", "deer", "elk", "moose", "reindeer", "white-tailed deer",
"raccoon", "skunk", "whale", "tasmanian devil", "badger", "ermine", "ferret", "marten", "mink", "otter", "weasel", "wolverine", "squirrel",
"white rabbit", "pegasus", "goat", "mouse", "cattle", "coyote", "owl", "charizard", "bandicoot", "giant panda", "panda", "bat", "ocelot", "panther",
"cheetah", "cougar", "puma", "pony", "camel", "llama", "penguin", "pigeon", "roadrunner", "seagull", "flamingo", "duck", 
    ],

    "pose": [
        "flexing", "double biceps pose", "side chest pose", "front lat spread pose", "back double biceps pose", "side biceps pose",
"adonis pose", "arms akimbo", "crossed arms", "hands on hips", "hands behind head", "one hand behind head", "torso twist", "power stance",
"v-shape", "clenched fists", "palm open", "fist pump", "victory pose", "comic shrug", "superhero pose", "front double biceps pose",
"front lat spread pose", "side chest pose", "side triceps pose", "rear double biceps pose", "rear lat spread pose",
"most muscular pose", "abdominal and thigh pose", "hands-on hips pose",
    ],

    "swimwear": [
        "swimming trunks", "swim shorts", "swim briefs", "speedo", "boxer swim shorts", "bermuda swim shorts", "swimming thong",
"swim g-string", "board shorts", "jammers", "tight swim shorts", "loose swim shorts", "short swim trunks",
    ],

    "headwear": [
        "baseball cap", "backwards headwear", "backwards baseball cap", "snapback cap", "backwards snapback cap", "trucker hat", "backwards trucker hat",
"fitted cap", "beanie hat", "headband",
"bandana",
    ],

    "setting:private": [
        "bedroom", "bathroom", "shower", "bathtub"
"locker room", "changing room", "dressing room", 
"communal shower", "forest", "tent", "poolside",
"sauna", "hot tub", "jacuzzi", 
"gym locker room", "living room", "beach",
    ],

    "setting:semi": [
        "bedroom", "living room", "locker room", "beach", "tent", "beach",
    ],
    
    "setting:public": [
        "living room", "library", "classroom", "train station", "subway station", "bus stop", "parking lot", "street", "alleyway",
"rooftop", "city street", "shopping mall", "supermarket", "gas station", "restaurant", "cafe", "bar",
    ],

    "setting:party": [
        "nightclub", "bar", "club", "dance floor", "music festival",
"rooftop party", "late-night street party", "party", "summer party", "beach party", "dance club", "strip club", "rave club", "bar counter",
    ],
    
    "setting:water": [
        "beach", "seashore", "lake shore", "riverbank", "swimming pool", "private pool", "poolside",
"deck by the pool", "beach lounge", "hot tub", "jacuzzi",
    ],
    
    "setting:underwater": [
        "pool", "jacuzzi", "hot tub", "spa pool", "sea", "lake", "river",
    ],

    "hair": [
        "bald", "short hair", "medium hair", "long hair", "man bun", "ponytail", "crew cut", "undercut", "side part", "spiky hair", "messy hair",
"curly hair", "wavy hair", "straight hair",
    ],

    "accessories": [
        "earrings", "stud earrings", "diamond stud earrings", "hoop earrings", "small hoop earrings", "septum piercing", "eyebrow piercing",
"chain necklace", "beaded necklace", "pendant necklace", "sunglasses", "aviator sunglasses", "wayfarer sunglasses", "sports sunglasses",
"headphones", "earphones",
    ],

    "hairy:face": [
        "beard", "goatee", "soul patch", "sideburns", "stubble", "full beard", "short beard", "long beard", "chin strap", "van dyke", "mutton chops", "scruffy beard", "designer stubble",
    ],

    "hairy:body": [
        "chest hair", "armpit hair", "happy trail", "pubic hair",
"chest hair + armpit hair", "chest hair + happy trail", "chest hair + pubic hair", "armpit hair + happy trail", "armpit hair + pubic hair",
"happy trail + pubic hair", "chest hair + armpit hair + happy trail", "chest hair + armpit hair + pubic hair", "chest hair + happy trail + pubic hair",
"armpit hair + happy trail + pubic hair", "chest hair + armpit hair + happy trail + pubic hair",
    ],

    "body": [
        "very skinny", "lean", "fit", "toned", "athletic", "muscular", "broad-shouldered", "stocky", "solid", "heavily muscled", "chubby", "big and muscular", "obese",
    ],

    "underwear": [
        "boxer briefs", "briefs", "boxers", "trunks", "thong", "g-string", "string thong", "jockstrap", "low-rise briefs", "mid-rise briefs",
"high-rise briefs", "athletic briefs",
    ],

    "sex:anal": [
        "mating press", "chair position", "cowgirl position", "doggystyle", "leg glider position", "mastery position",
"Spoon Position", "reverse missionary position", "missionary position", "reverse cowgirl position",
"anvil Position", "guard Position", "stand and carry Position", "table lotus position"
    ],

    "sex:oral": [
        "69 position", "kneeling oral position", "sideways oral", "penis lick", "tongue in foreskin", "deepthroath", "tongue out blowjob", "irrumatio", 
"ball lick", "kneeling and blow position", "lying and blow position", "north pole position", "sit and blow position",
    ],

    "topwear:under": [
        "tank top", "tank top", "loose tank top, extended armholes, draped sleeveless shirt", "oversized muscle tank, baggy sleeveless top, open armholes",
"t-shirt", "crew-neck t-shirt", "v-neck t-shirt", "sleeveless shirt", "cotton undershirt",
    ],

    "topwear:casual": [
        "long-sleeve t-shirt", "crew-neck long-sleeve shirt", "v-neck long-sleeve shirt", "henley long-sleeve", "hoodie",
"sweatshirt", "zip-up hoodie", "pullover sweater", "thermal shirt", "flannel shirt",
    ],

    "topwear:formal": [
        "button-up long-sleeve shirt", "collared shirt", "dress shirt", "oxford shirt", "polo shirt", "unbuttoned shirt, flannel shirt,"
    ],

    "topwear:party": [
        "mesh tank top", "mesh shirt", "harness", "cage top", "leather vest", "rubber shirt", "tight tank top", "open shirt", "crop top",
"fishnet top", "bondage harness", "strappy top", "latex shirt", "transparent shirt", "studded vest", "festival top", "topless", "topless",
    ],
    
    "bottomwear:formal": [
        "jeans", "chinos", "khaki pants", "linen trousers", "cotton trousers", "dress trousers", "slacks", "wool trousers", "formal pants",
    ],

    "bottomwear:casual": [
        "joggers", "sweatpants", "track pants", "gym shorts", "shorts", "linen shorts", "bermuda shorts", "cargo pants", "basketball shorts", "training pants",
    ],

    "bottomwear:party": [
        "camo pants", "leather pants", "rubber pants", "cargo pants", "denim pants", "tight jeans", "distressed jeans", "ripped jeans", "patent leather pants",
"bondage pants", "festival trousers", "shiny pants", "studded pants", "paint-splatter pants", "metallic pants", "faux leather pants", "tight cargo pants",
"plaid pants",
    ],

    "light": [
        "daylight", "bright light", "soft light", "dim light", "night", "ambient light", "diffuse light", "even lighting", "natural light", "low light",
"dappled light", "soft shadows", "subtle highlights", "warm light", "cool light" "dusty light", "hazy light", "soft glow", "subtle lens flare",
"misty light", "warm haze", "cool haze",
    ],

    "perspective": [
        "front view", "side view", "three-quarter view", "top-down view", "worm's-eye view", "high-angle view", "low-angle view",
"overhead view", "diagonal view", "isometric view", "dutch angle", "fisheye perspective", "foreshortening",
    ],

    "effect": [
        "soft focus", "slight blur", "bokeh", "grainy texture", "film grain", "vignette", "subtle glow", "soft haze", "hazy atmosphere",
"light diffusion", "slight overexposure", "subtle underexposure", "warm tint", "cool tint", "muted colors", "slightly desaturated",
"painterly effect", "sketch effect", "cartoonish style", "comic style"
    ],

    "face": [
        "flirty", "seductive", "playful flirt", "teasing look", "sultry gaze", "wink flirt", "one eye closed", "smoldering look", "cheeky flirt", "suggestive grin",
"mischievous smirk", "playful smirk", "tongue out flirty", "open mouth flirt", "eye contact flirt", "raised eyebrow flirt", "smiling", "grinning", "laughing", "happy",
"joyful", "playful", "quirky", "cheerful", "friendly smile", "tongue out", "winking", "raised eyebrows", "blushing", "embarrassed", "shy", "bashful", "surprised",
"curious", "amused", "satisfied smile", "soft smile", "teasing smile", "cheeky grin", "slight smirk", "mischievous grin", "open mouth smile", "excited", "delighted",
"flirty", "seductive", "playful flirt", "teasing look", "sultry gaze", "wink flirt",
    ],

    "color": [
        "red", "scarlet", "crimson", "pink", "magenta", "fuchsia", "orange", "tangerine", "yellow", "bright yellow", "lemon yellow color", "gold color", "green", "lime green color",
"emerald", "olive", "teal", "turquoise", "cyan", "blue", "sky blue", "royal blue", "navy blue", "indigo", "purple", "violet", "lavender", "lilac", "brown", "tan",
"beige", "bronze", "maroon", "burgundy", "white", "ivory", "black", "grey", "dark-grey", "light-grey", "silver", "reflecting chrome", "copper", "mint", "jade", 
"peach color", "camo", "camouflage", "red-violet", "blue-green", "turquoise-blue", "yellow-orange", "pink-purple", "mustard color",
    ],
}

def apply_scratchpad(prompt: str) -> str:
    """Korvaa [tagi] scratchpadin vaihtoehdoilla satunnaisesti.
       Jos tagia ei löydy, poistetaan hakasulkeet kokonaan.
    """
    matches = re.findall(r"\[(.+?)\]", prompt)
    for tag in matches:
        if tag in scratchpad:
            choice = random.choice(scratchpad[tag])
            # korvataan ilman hakasulkeita
            prompt = prompt.replace(f"[{tag}]", choice, 1)
        else:
            # poistetaan tuntemattomat tagit
            prompt = prompt.replace(f"[{tag}]", "", 1)
    return prompt.strip()


resolutions = [
    (592, 1768), (608, 1696), (640, 1624), (672, 1544),
    (720, 1440), (744, 1392), (784, 1320), (832, 1248),
    (864, 1208), (880, 1184), (912, 1144), (976, 1072),
    (1024, 1024), (1072, 912), (1144, 912), (1184, 880),
    (1208, 864), (1248, 832), (1320, 784), (1392, 744),
    (1440, 720), (1544, 672), (1624, 640), (1696, 608),
    (1768, 592),
    (720, 1440), (744, 1392), (784, 1320), (832, 1248),
    (864, 1208), (880, 1184), (912, 1144), (976, 1072),
    (1024, 1024), (1072, 912), (1144, 912), (1184, 880),
    (1208, 864), (1248, 832), (1320, 784),
    (880, 1184), (912, 1144), (976, 1072),
    (1024, 1024), (1072, 912),
    (912, 1144), (976, 1072),
    (1024, 1024),
]


@spaces.GPU
def generate(prompt, negative_prompt, width, height, random_res, selected_style):
    if random_res:
        width, height = random.choice(resolutions)

    # Expandataan scratchpad-tagit
    prompt = apply_scratchpad(prompt)


    # Ladataan pysyvät + tyyli-LoRA yhteen
    all_loras = {
        "kkntr/lora1": 0.2,
        "kkntr/lora2": 0.3,
        "kkntr/lora12": 0.4,
        "kkntr/lora10": 0.1,
    }

    pipe.load_lora_weights(all_loras)
    pipe.fuse_lora()

    quality_suffix = (
        ", (anthro), (male), (best quality), high resolution, (ultra detailed),"
        " 8k, (highly detailed), (masterpiece), (sharp focus), (intricate details),"
        " (perfect anatomy), (depth of field), high dynamic range,"
        " (smooth shading), (realistic shadows), (rich texture),"
        " professional color grading, (fine details), ultra sharp,"
        " HDR, (anatomically correct), painterly realism, 20r0j_2, civifur"
    )

    full_prompt = prompt + quality_suffix

    image = pipe(
        prompt=full_prompt,
        negative_prompt=negative_prompt,
        width=width,
        height=height,
        num_inference_steps=40,
        guidance_scale=7.5
    ).images[0]

    return image, full_prompt

gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="Prompt"),
        gr.Textbox(
            label="Negative Prompt",
            value="(worst quality), (low quality), (blurry), (bad anatomy), (fused fingers), (extra fingers), (missing fingers), (deformed), (jpeg artifacts), (watermark), (signature), text, error, out of frame, bad hands, bad feet, bad proportions, fused limbs, cropped, username, logo, bad perspective, lowres, bad artist, bad composition, distorted, mutated, bad shadow, multiple limbs, extra arms, extra legs, bad background, tilted, bad lighting, bad face, ugly, bad eyes, disfigured, multiple pictures"
        ),
        gr.Slider(minimum=512, maximum=1920, step=64, value=768, label="Width"),
        gr.Slider(minimum=512, maximum=1920, step=64, value=768, label="Height"),
        gr.Checkbox(label="Random resolution mode (override sliders)", value=True),
    ],
    outputs=[
        gr.Image(label="Generated Image"),
        gr.Textbox(label="Used Prompt", lines=4)
    ],
    title="FURRYSTYLE",
    description="animal,   body,   sex:anal/oral,   pose,   hairy:face/body,   "
    "topwear:under/casual/formal/party,   bottomwear:casual/formal/party,   " 
    "underwear,   swimwear,   headwear,   hair,   accessories,   face,   light,   perspective,   "
    "effect,   setting:private/semi/public/party/water/underwater,   color"
).launch()