from __future__ import annotations

import argparse
import importlib
import importlib.util
import shutil
import sys
import time
import traceback
from dataclasses import dataclass
from pathlib import Path
from types import ModuleType

import settings
from captioning import clean_tag, format_caption_for_training
from image_prep import prepare_matched_datasets, reset_dir


@dataclass
class TrainOptions:
    mode: str
    model_type: str
    base_model: str
    anima_qwen3: str
    anima_vae: str
    anima_llm_adapter: str
    output_name: str
    target_caption: str
    pair_caption: str
    lora_output_dir: str
    max_train_steps: int
    learning_rate: float
    train_batch_size: int
    network_dim: int
    merge_strength: float
    network_alpha: float | None = None
    speed_preset: str = ""
    low_vram: bool = False
    gradient_checkpointing: bool = True
    persist_output_dir: bool = True
    save_step_variants: bool = False
    grayscale_reference_for_color: bool = False


def ensure_dirs() -> None:
    settings.RUNS_DIR.mkdir(parents=True, exist_ok=True)
    settings.OUTPUT_DIR.mkdir(parents=True, exist_ok=True)


def scan_files(model_dir: str | Path | None, fallback: Path, extensions: tuple[str, ...]) -> list[str]:
    files: list[str] = []
    directory = settings.normalize_user_path(model_dir, fallback) if model_dir else fallback
    if not directory.exists():
        return files
    for pattern in tuple(f"*{extension}" for extension in extensions):
        for path in sorted(directory.rglob(pattern)):
            files.append(str(path))
    return files


def scan_models(model_dir: str | Path | None = None) -> list[str]:
    return scan_files(model_dir, settings.get_sdxl_model_dir(), (".safetensors", ".ckpt"))


def scan_anima_files(model_dir: str | Path | None = None) -> list[str]:
    return scan_files(model_dir, settings.get_anima_model_dir(), (".safetensors", ".pth"))


ANIMA_DIT_REQUIRED_KEYS = (
    "x_embedder.proj.1.weight",
    "t_embedder.1.linear_1.weight",
    "t_embedder.1.linear_2.weight",
    "llm_adapter.embed.weight",
    "llm_adapter.blocks.0.norm_self_attn.weight",
    "llm_adapter.blocks.0.self_attn.q_proj.weight",
)

SDXL_HINT_KEYS = (
    "model.diffusion_model.input_blocks.0.0.weight",
    "conditioner.embedders.0.model.text_model.embeddings.token_embedding.weight",
    "first_stage_model.encoder.conv_in.weight",
)


def _has_key(keys: set[str], suffix: str) -> bool:
    if suffix in keys or f"net.{suffix}" in keys:
        return True
    return any(key.endswith(f".{suffix}") for key in keys)


def read_safetensors_keys(path: str | Path) -> set[str]:
    from safetensors import safe_open

    model_path = Path(path)
    with safe_open(str(model_path), framework="pt") as handle:
        return set(handle.keys())


def is_anima_dit_checkpoint(path: str | Path) -> bool:
    try:
        ok, _ = validate_anima_dit_checkpoint(path, quiet=True)
        return ok
    except Exception:
        return False


def validate_anima_dit_checkpoint(path: str | Path, quiet: bool = False) -> tuple[bool, str]:
    model_path = Path(path)
    if not model_path.exists():
        return False, f"Animaモデルが見つかりません。\n{model_path}"

    if model_path.suffix.lower() != ".safetensors":
        return (
            False,
            "Anima学習には Anima 本体の .safetensors が必要です。\n"
            f"選択中: {model_path}",
        )

    try:
        keys = read_safetensors_keys(model_path)
    except Exception as exc:
        return (
            False,
            "Animaモデルの中身を確認できませんでした。\n"
            f"選択中: {model_path}\n\n{exc}",
        )

    missing = [key for key in ANIMA_DIT_REQUIRED_KEYS if not _has_key(keys, key)]
    if not missing:
        return True, "Anima DiTモデルとして確認できました。"

    looks_like_sdxl = any(_has_key(keys, key) for key in SDXL_HINT_KEYS)
    model_kind = "SDXL/Illustrious系チェックポイントに見えます。\n" if looks_like_sdxl else ""
    missing_text = ", ".join(missing[:3])
    if quiet:
        return False, missing_text
    return (
        False,
        "選択中のAnimaモデルは、この学習方式で使うAnima DiT形式ではありません。\n"
        f"{model_kind}"
        f"選択中: {model_path}\n\n"
        "必要なDiTキーが見つかりません。\n"
        f"不足例: {missing_text}\n\n"
        "modelsフォルダ内の StableDiffusion フォルダに、Anima本体のDiT .safetensorsを入れてから、"
        "モデルを読み込み直してください。名前にAnimaが入ったSDXL系モデルはAnima学習には使えません。",
    )


def check_environment() -> str:
    messages: list[str] = []
    if settings.SD_SCRIPTS_DIR.exists():
        messages.append(f"sd-scripts: OK ({settings.SD_SCRIPTS_DIR})")
    else:
        messages.append("sd-scripts: 未セットアップです。Copita/install.ps1 を実行してください。")

    for module in ("torch", "torchvision", "accelerate", "gradio", "PIL", "toml", "safetensors"):
        if importlib.util.find_spec(module) is None:
            messages.append(f"{module}: 未インストール")
        else:
            messages.append(f"{module}: OK")

    try:
        import torch

        if torch.cuda.is_available():
            messages.append(f"GPU: OK ({torch.cuda.get_device_name(0)})")
            messages.append(f"PyTorch CUDA: {torch.__version__} / CUDA {torch.version.cuda}")
        else:
            cuda_build = torch.version.cuda or "なし"
            messages.append(f"GPU: 未使用（PyTorch {torch.__version__}, CUDA {cuda_build}）")
        messages.append("LoRA学習にはGPU版PyTorchが必要です。必要に応じて setup_first_time.bat をもう一度実行してください。")
    except Exception as exc:
        messages.append(f"GPU確認: 失敗 ({exc})")

    models = scan_models(settings.get_sdxl_model_dir())
    anima_models = scan_models(settings.get_anima_model_dir())
    messages.append(f"modelsフォルダ: {settings.get_models_root_dir()}")
    messages.append(f"StableDiffusion: {settings.get_sdxl_model_dir()} ({len(models)}件)")
    messages.append(f"TextEncoders: {settings.get_anima_text_encoder_dir()}")
    messages.append(f"VAE: {settings.get_anima_vae_dir()}")
    messages.append(f"Anima候補: {len(anima_models)}件")
    if (settings.SD_SCRIPTS_DIR / "anima_train_network.py").exists():
        messages.append("Anima学習: OK")
    else:
        messages.append("Anima学習: sd-scriptsを最新版に更新してください。setup_first_time.bat をもう一度実行します。")
    try:
        messages.append(f"LoRA保存先: {settings.get_lora_output_dir()}")
    except Exception as exc:
        messages.append(f"LoRA保存先: 設定を読めませんでした ({exc})")
    return "\n".join(messages)


def load_module_from_path(module_name: str, path: Path) -> ModuleType:
    spec = importlib.util.spec_from_file_location(module_name, path)
    if spec is None or spec.loader is None:
        raise RuntimeError(f"Pythonモジュールを読み込めません: {path}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def add_sd_scripts_to_path() -> None:
    sd_scripts_dir = settings.SD_SCRIPTS_DIR
    if not sd_scripts_dir.exists():
        raise RuntimeError(
            "sd-scripts が見つかりません。まず Copita/install.ps1 を実行してください。"
        )

    for sub in ("", "networks", "library", "tools"):
        path = str(sd_scripts_dir / sub) if sub else str(sd_scripts_dir)
        if path not in sys.path:
            sys.path.append(path)


def import_sdxl_scripts():
    add_sd_scripts_to_path()
    sd_scripts_dir = settings.SD_SCRIPTS_DIR

    import sdxl_train_network

    merge_path = sd_scripts_dir / "networks" / "sdxl_merge_lora.py"
    resize_path = sd_scripts_dir / "networks" / "resize_lora.py"
    if not merge_path.exists() or not resize_path.exists():
        raise RuntimeError("sd-scripts の networks/sdxl_merge_lora.py または resize_lora.py が見つかりません。")

    merge = load_module_from_path("sdxl_merge_lora_tool", merge_path)
    resize = load_module_from_path("resize_lora_tool", resize_path)

    return sdxl_train_network, merge, resize


def import_anima_scripts():
    add_sd_scripts_to_path()
    sd_scripts_dir = settings.SD_SCRIPTS_DIR
    anima_path = sd_scripts_dir / "anima_train_network.py"
    merge_path = sd_scripts_dir / "networks" / "merge_lora.py"
    resize_path = sd_scripts_dir / "networks" / "resize_lora.py"

    if not anima_path.exists():
        raise RuntimeError(
            "このsd-scriptsには Anima 学習スクリプトがありません。\n"
            "setup_first_time.bat をもう一度実行して、sd-scriptsを最新版に更新してください。"
        )
    if not merge_path.exists() or not resize_path.exists():
        raise RuntimeError("sd-scripts の networks/merge_lora.py または resize_lora.py が見つかりません。")

    import anima_train_network

    merge = load_module_from_path("generic_merge_lora_tool", merge_path)
    resize = load_module_from_path("generic_resize_lora_tool", resize_path)

    return anima_train_network, merge, resize


def training_args_util(train_module):
    util = getattr(train_module, "train_util", None)
    if util is not None:
        return util
    util = getattr(train_module, "args_util", None)
    if util is not None:
        return util

    import library.args as args_util

    return args_util


def reset_sd_strategies() -> None:
    try:
        from library.strategy_base import (
            LatentsCachingStrategy,
            TextEncoderOutputsCachingStrategy,
            TextEncodingStrategy,
            TokenizeStrategy,
        )
    except Exception:
        return

    for cls in (
        TokenizeStrategy,
        TextEncodingStrategy,
        TextEncoderOutputsCachingStrategy,
        LatentsCachingStrategy,
    ):
        cls._strategy = None


def prepare_tqdm_progress_bars() -> None:
    try:
        import tqdm as tqdm_module
    except Exception:
        return

    base_tqdm = getattr(tqdm_module, "_copita_base_tqdm", None)
    if base_tqdm is None:
        base_tqdm = tqdm_module.tqdm
        setattr(tqdm_module, "_copita_base_tqdm", base_tqdm)

    def ui_tqdm(*args, **kwargs):
        kwargs["disable"] = False
        kwargs.setdefault("ascii", True)
        kwargs.setdefault("mininterval", 1.0)
        return base_tqdm(*args, **kwargs)

    tqdm_module.tqdm = ui_tqdm
    for module_name in (
        "train_network",
        "sdxl_train_network",
        "anima_train_network",
        "library.train_util",
        "library.config_util",
    ):
        module = sys.modules.get(module_name)
        if module is not None and hasattr(module, "tqdm"):
            setattr(module, "tqdm", ui_tqdm)


def effective_network_alpha(options: TrainOptions, default: float) -> float:
    value = options.network_alpha
    if value is None:
        return float(default)
    try:
        alpha = float(value)
    except Exception:
        return float(default)
    return alpha if alpha > 0 else float(default)


def anima_experiment_tuning(network_dim: int, network_alpha: float, learning_rate: float) -> dict[str, float | int]:
    dim = int(network_dim)
    lr = float(learning_rate)
    return {
        "network_alpha": float(network_alpha),
        "cross_dim": max(4, dim // 2),
        "self_lr": lr,
        "cross_lr": lr * 0.7,
        "mlp_lr": lr * 0.8,
    }


def intermediate_checkpoint_steps(options: TrainOptions) -> list[int]:
    if not options.save_step_variants:
        return []
    max_steps = int(options.max_train_steps)
    if max_steps <= 300:
        return []
    first_step = max(100, max_steps - 200)
    return list(range(first_step, max_steps, 100))


def step_checkpoint_path(output_dir: Path, output_name: str, step: int) -> Path:
    return output_dir / f"{output_name}-step{step:08d}.safetensors"


def planned_output_paths(output_dir: Path, output_name: str, options: TrainOptions) -> list[Path]:
    paths = [output_dir / f"{output_name}.safetensors"]
    paths.extend(output_dir / f"{output_name}_step{step}.safetensors" for step in intermediate_checkpoint_steps(options))
    return paths


def apply_intermediate_checkpoint_settings(args_dict: dict, options: TrainOptions) -> None:
    if intermediate_checkpoint_steps(options):
        args_dict["save_every_n_steps"] = 100


def use_low_vram_mode(options: TrainOptions) -> bool:
    return bool(options.low_vram)


def train_one_lora(
    sdxl_train_network,
    base_model: str,
    train_data_dir: Path,
    output_dir: Path,
    output_name: str,
    options: TrainOptions,
) -> Path:
    reset_sd_strategies()
    output_dir.mkdir(parents=True, exist_ok=True)

    args_dict = {
        "pretrained_model_name_or_path": base_model,
        "train_data_dir": str(train_data_dir),
        "output_dir": str(output_dir),
        "output_name": output_name,
        "max_train_steps": int(options.max_train_steps),
        "network_module": "networks.lora",
        "xformers": False,
        "sdpa": True,
        "gradient_checkpointing": True,
        "persistent_data_loader_workers": False,
        "max_data_loader_n_workers": 0,
        "enable_bucket": True,
        "save_model_as": "safetensors",
        "lr_scheduler_num_cycles": 4,
        "learning_rate": float(options.learning_rate),
        "resolution": "1024,1024",
        "train_batch_size": int(options.train_batch_size),
        "network_dim": int(options.network_dim),
        "network_alpha": effective_network_alpha(options, int(options.network_dim)),
        "optimizer_type": "AdamW8bit",
        "mixed_precision": "fp16",
        "save_precision": "fp16",
        "lr_scheduler": "constant",
        "bucket_no_upscale": True,
        "min_bucket_reso": 64,
        "max_bucket_reso": 1024,
        "caption_extension": ".txt",
        "seed": settings.DEFAULT_TRAIN_SETTINGS["seed"],
        "network_train_unet_only": True,
        "no_half_vae": True,
        "cache_latents": True,
        "cache_latents_to_disk": True,
        "cache_text_encoder_outputs": True,
        "cache_text_encoder_outputs_to_disk": True,
        "fp8_base": True,
        "console_log_simple": True,
        "lowram": use_low_vram_mode(options),
    }
    apply_intermediate_checkpoint_settings(args_dict, options)

    parser = sdxl_train_network.setup_parser()
    args = parser.parse_args([])
    args_util = training_args_util(sdxl_train_network)
    args_util.verify_command_line_training_args(args)
    args = args_util.read_config_from_file(args, parser)
    for key, value in args_dict.items():
        setattr(args, key, value)

    prepare_tqdm_progress_bars()
    trainer = sdxl_train_network.SdxlNetworkTrainer()
    trainer.train(args)

    lora_path = output_dir / f"{output_name}.safetensors"
    if not lora_path.exists():
        raise RuntimeError(f"学習は終了しましたが、出力LoRAが見つかりません: {lora_path}")
    return lora_path


def train_one_anima_lora(
    anima_train_network,
    base_model: str,
    train_data_dir: Path,
    output_dir: Path,
    output_name: str,
    options: TrainOptions,
) -> Path:
    reset_sd_strategies()
    output_dir.mkdir(parents=True, exist_ok=True)
    use_fast_plus = options.speed_preset in {
        "通常学習セット",
        "詳細学習セット",
        "通常（16GB以上）",
        "高速（40分目標）",
        "高品質",
        "高速（16GB以上）",
        "高速+（16GB実験）",
        "高速++（16GB実験2）",
    }

    network_dim = int(options.network_dim)
    network_alpha = effective_network_alpha(options, settings.DEFAULT_ANIMA_TRAIN_SETTINGS["network_alpha"])
    tuning = anima_experiment_tuning(network_dim, network_alpha, float(options.learning_rate))
    cross_dim = int(tuning["cross_dim"])
    network_args = [
        f"network_reg_dims=.*self_attn.*={network_dim},.*cross_attn.*={cross_dim},.*mlp.*={network_dim}",
        f"network_reg_lrs=.*self_attn.*={float(tuning['self_lr'])},.*cross_attn.*={float(tuning['cross_lr'])},.*mlp.*={float(tuning['mlp_lr'])}",
    ]

    args_dict = {
        "pretrained_model_name_or_path": base_model,
        "qwen3": options.anima_qwen3,
        "vae": options.anima_vae,
        "llm_adapter_path": options.anima_llm_adapter or None,
        "train_data_dir": str(train_data_dir),
        "output_dir": str(output_dir),
        "output_name": output_name,
        "max_train_steps": int(options.max_train_steps),
        "network_module": "networks.lora_anima",
        "gradient_checkpointing": True,
        "persistent_data_loader_workers": False,
        "max_data_loader_n_workers": 0,
        "enable_bucket": True,
        "save_model_as": "safetensors",
        "learning_rate": float(options.learning_rate),
        "resolution": "1024,1024",
        "train_batch_size": int(options.train_batch_size),
        "network_dim": network_dim,
        "network_alpha": float(tuning["network_alpha"]),
        "optimizer_type": "AdamW8bit",
        "mixed_precision": "bf16",
        "save_precision": "bf16",
        "lr_scheduler": "constant",
        "timestep_sampling": "sigmoid",
        "discrete_flow_shift": 1.0,
        "network_train_unet_only": True,
        "cache_latents": True,
        "cache_latents_to_disk": True,
        "cache_text_encoder_outputs": True,
        "cache_text_encoder_outputs_to_disk": True,
        "caption_extension": ".txt",
        "seed": settings.DEFAULT_ANIMA_TRAIN_SETTINGS["seed"],
        "console_log_simple": True,
        "lowram": use_low_vram_mode(options),
        "attn_mode": "torch",
        "split_attn": not use_fast_plus,
        "vae_chunk_size": 64,
        "vae_disable_cache": True,
        "network_args": network_args,
    }
    apply_intermediate_checkpoint_settings(args_dict, options)

    parser = anima_train_network.setup_parser()
    args = parser.parse_args([])
    args_util = training_args_util(anima_train_network)
    args_util.verify_command_line_training_args(args)
    args = args_util.read_config_from_file(args, parser)
    for key, value in args_dict.items():
        setattr(args, key, value)

    prepare_tqdm_progress_bars()
    trainer = anima_train_network.AnimaNetworkTrainer()
    trainer.train(args)

    lora_path = output_dir / f"{output_name}.safetensors"
    if not lora_path.exists():
        raise RuntimeError(f"学習は終了しましたが、出力LoRAが見つかりません: {lora_path}")
    return lora_path


def merge_difference_lora(merge_module, resize_module, positive_lora: Path, negative_lora: Path, output_path: Path, options: TrainOptions) -> Path:
    temp_merge = output_path.with_name(output_path.stem + "_merged_raw.safetensors")
    merge_args = argparse.Namespace(
        save_precision="bf16",
        precision="float",
        sd_model=None,
        save_to=str(temp_merge),
        models=[str(positive_lora), str(negative_lora)],
        no_metadata=False,
        ratios=[float(options.merge_strength), -float(options.merge_strength)],
        concat=True,
        shuffle=True,
        lbws=[],
    )
    merge_module.merge(merge_args)

    resize_args = argparse.Namespace(
        save_precision="bf16",
        new_rank=int(options.network_dim),
        new_conv_rank=int(options.network_dim),
        save_to=str(output_path),
        model=str(temp_merge),
        device="cuda",
        verbose=True,
        dynamic_param=None,
        dynamic_method=None,
        svd_lowrank_niter=2,
        lbws=[],
    )
    resize_module.resize(resize_args)
    rewrite_lora_metadata(
        output_path,
        "networks.lora",
        effective_network_alpha(options, int(options.network_dim)),
        int(options.network_dim),
    )
    try:
        temp_merge.unlink(missing_ok=True)
    except Exception:
        pass
    return output_path


def rewrite_lora_metadata(
    lora_path: Path,
    network_module: str,
    network_alpha: float | None = None,
    network_dim: int | None = None,
) -> None:
    try:
        from safetensors import safe_open
        from safetensors.torch import load_file, save_file
    except Exception:
        return

    try:
        with safe_open(str(lora_path), framework="pt") as handle:
            metadata = dict(handle.metadata() or {})
        tensors = load_file(str(lora_path), device="cpu")
        metadata["ss_network_module"] = network_module
        if network_dim is not None:
            metadata["ss_network_dim"] = str(int(network_dim))
        if network_alpha is not None:
            metadata["ss_network_alpha"] = str(float(network_alpha))
            for key, value in list(tensors.items()):
                if key.endswith(".alpha"):
                    tensors[key] = value.new_tensor(float(network_alpha))
        temp_path = lora_path.with_name(f"{lora_path.stem}.metadata_tmp{lora_path.suffix}")
        save_file(tensors, str(temp_path), metadata=metadata)
        temp_path.replace(lora_path)
    except Exception as exc:
        try:
            temp_path.unlink(missing_ok=True)  # type: ignore[name-defined]
        except Exception:
            pass
        print(f"LoRAメタデータの更新に失敗しました: {lora_path} ({exc})", flush=True)
        return


def merge_difference_lora_generic(
    merge_module,
    resize_module,
    positive_lora: Path,
    negative_lora: Path,
    output_path: Path,
    options: TrainOptions,
    network_module: str,
) -> Path:
    temp_merge = output_path.with_name(output_path.stem + "_merged_raw.safetensors")
    merge_args = argparse.Namespace(
        save_precision="bf16",
        precision="float",
        sd_model=None,
        save_to=str(temp_merge),
        models=[str(positive_lora), str(negative_lora)],
        no_metadata=False,
        ratios=[float(options.merge_strength), -float(options.merge_strength)],
        concat=True,
        shuffle=True,
        lbws=[],
        v2=False,
    )
    merge_module.merge(merge_args)
    rewrite_lora_metadata(temp_merge, network_module)

    resize_args = argparse.Namespace(
        save_precision="bf16",
        new_rank=int(options.network_dim),
        new_conv_rank=int(options.network_dim),
        save_to=str(output_path),
        model=str(temp_merge),
        device="cuda",
        verbose=True,
        dynamic_param=None,
        dynamic_method=None,
        svd_lowrank_niter=2,
        lbws=[],
    )
    resize_module.resize(resize_args)
    rewrite_lora_metadata(output_path, network_module)
    try:
        temp_merge.unlink(missing_ok=True)
    except Exception:
        pass
    return output_path


def publish_lora_copy(final_path: Path, run_dir: Path) -> None:
    if final_path.parent.resolve() != settings.OUTPUT_DIR.resolve():
        shutil.copy2(final_path, settings.OUTPUT_DIR / final_path.name)
    shutil.copy2(final_path, run_dir / final_path.name)


def merge_checkpoint_differences(
    merge_module,
    resize_module,
    raw_lora_dir: Path,
    output_dir: Path,
    output_name: str,
    options: TrainOptions,
    model_type: str,
    run_dir: Path,
) -> list[Path]:
    output_paths: list[Path] = []
    network_module = "networks.lora_anima" if model_type == "ANIMA" else "networks.lora"
    for step in intermediate_checkpoint_steps(options):
        target_lora = step_checkpoint_path(raw_lora_dir, f"{output_name}_target", step)
        reference_lora = step_checkpoint_path(raw_lora_dir, f"{output_name}_reference", step)
        if not target_lora.exists() or not reference_lora.exists():
            print(f"途中LoRA {step} step が見つからないためスキップします。", flush=True)
            continue
        checkpoint_output = output_dir / f"{output_name}_step{step}.safetensors"
        print(f"途中差分LoRAを作成中: {step} step", flush=True)
        if model_type == "ANIMA":
            merge_difference_lora_generic(
                merge_module,
                resize_module,
                target_lora,
                reference_lora,
                checkpoint_output,
                options,
                network_module,
            )
        else:
            merge_difference_lora(merge_module, resize_module, target_lora, reference_lora, checkpoint_output, options)
        publish_lora_copy(checkpoint_output, run_dir)
        output_paths.append(checkpoint_output)
    return output_paths


def safe_output_name(name: str) -> str:
    cleaned = clean_tag(name)
    if cleaned:
        return cleaned
    return "style_diff_lora"


def run_training(target_image: str, pair_image: str | None, options: TrainOptions) -> tuple[str, str | None]:
    ensure_dirs()
    if not target_image:
        return "学習したい画風の画像を選んでください。", None

    is_pair_mode = options.mode == "pair"
    if is_pair_mode and not pair_image:
        return "画風だけが違う比較画像も選んでください。", None
    if not is_pair_mode and not pair_image:
        return "画像生成AIで作った比較画像を読み込んでから学習を開始してください。", None

    model_type = (options.model_type or "SDXL").upper()
    if not options.base_model or not Path(options.base_model).exists():
        if model_type == "ANIMA":
            return "Animaモデルが見つかりません。Anima設定で本体の.safetensorsを選んでください。", None
        return "SDXLモデルが見つかりません。学習設定で、普段使っているモデルフォルダを指定してからモデルを選んでください。", None
    if model_type == "ANIMA":
        is_valid_anima, anima_message = validate_anima_dit_checkpoint(options.base_model)
        if not is_valid_anima:
            return anima_message, None
        if not options.anima_qwen3 or not Path(options.anima_qwen3).exists():
            return "AnimaのQwen3テキストエンコーダーが見つかりません。Anima設定でQwen3を選んでください。", None
        if not options.anima_vae or not Path(options.anima_vae).exists():
            return "AnimaのVAEが見つかりません。Anima設定でVAEを選んでください。", None
        if options.anima_llm_adapter and not Path(options.anima_llm_adapter).exists():
            return "Animaの追加コンポーネントが見つかりません。空欄にするか、正しいファイルを選んでください。", None

    run_id = time.strftime("%Y%m%d_%H%M%S")
    output_name = safe_output_name(options.output_name)
    lora_output_dir = settings.normalize_user_path(options.lora_output_dir, settings.get_lora_output_dir())
    lora_output_dir.mkdir(parents=True, exist_ok=True)
    existing_outputs = [path for path in planned_output_paths(lora_output_dir, output_name, options) if path.exists()]
    if existing_outputs:
        existing_text = "\n".join(f"- {path}" for path in existing_outputs)
        return (
            "同じ名前のLoRAがすでに保存先にあります。上書きはしません。\n"
            "LoRA名を変えるか、保存先を変更してください。\n"
            f"{existing_text}",
            None,
        )
    if options.persist_output_dir:
        settings.save_user_settings({"lora_output_dir": str(lora_output_dir)})
    run_dir = settings.RUNS_DIR / f"{run_id}_{output_name}"
    reset_dir(run_dir)

    try:
        assert pair_image is not None
        raw_lora_dir = run_dir / "raw_loras"
        target_caption = format_caption_for_training(options.target_caption, model_type)
        pair_caption = format_caption_for_training(options.pair_caption or options.target_caption, model_type)
        target_data, reference_data = prepare_matched_datasets(
            target_image,
            pair_image,
            target_caption,
            pair_caption,
            run_dir,
            grayscale_pair=options.grayscale_reference_for_color,
        )

        if model_type == "ANIMA":
            anima_train_network, merge_module, resize_module = import_anima_scripts()
            print("工程 1/3: 学習画像LoRAを学習中", flush=True)
            target_lora = train_one_anima_lora(
                anima_train_network,
                options.base_model,
                target_data,
                raw_lora_dir,
                f"{output_name}_target",
                options,
            )
            print("工程 2/3: 比較画像LoRAを学習中", flush=True)
            negative_lora = train_one_anima_lora(
                anima_train_network,
                options.base_model,
                reference_data,
                raw_lora_dir,
                f"{output_name}_reference",
                options,
            )
        else:
            sdxl_train_network, merge_module, resize_module = import_sdxl_scripts()
            print("工程 1/3: 学習画像LoRAを学習中", flush=True)
            target_lora = train_one_lora(
                sdxl_train_network,
                options.base_model,
                target_data,
                raw_lora_dir,
                f"{output_name}_target",
                options,
            )
            print("工程 2/3: 比較画像LoRAを学習中", flush=True)
            negative_lora = train_one_lora(
                sdxl_train_network,
                options.base_model,
                reference_data,
                raw_lora_dir,
                f"{output_name}_reference",
                options,
            )
        final_path = lora_output_dir / f"{output_name}.safetensors"
        print("工程 3/3: 差分LoRAを作成中", flush=True)
        intermediate_paths = merge_checkpoint_differences(
            merge_module,
            resize_module,
            raw_lora_dir,
            lora_output_dir,
            output_name,
            options,
            model_type,
            run_dir,
        )
        if model_type == "ANIMA":
            merge_difference_lora_generic(merge_module, resize_module, target_lora, negative_lora, final_path, options, "networks.lora_anima")
        else:
            merge_difference_lora(merge_module, resize_module, target_lora, negative_lora, final_path, options)

        publish_lora_copy(final_path, run_dir)
        intermediate_text = ""
        if intermediate_paths:
            intermediate_text = "\n途中差分:\n" + "\n".join(f"- {path}" for path in intermediate_paths) + "\n"
        message = (
            "差分LoRAが完成しました。\n"
            f"出力先: {final_path}\n"
            f"{intermediate_text}"
            f"学習タイプ: {model_type}\n"
            "\n"
            "生成時はLoRA強度 0.6-0.8程度で試してください"
        )
        if options.persist_output_dir:
            settings.save_user_settings({"lora_output_dir": str(lora_output_dir), "last_lora_output_name": output_name})
        return message, str(final_path)
    except Exception:
        return "エラーが発生しました。\n\n" + traceback.format_exc(), None
