import torch
from nodes import LoraLoader, LoraLoaderModelOnly

class MultiLoRAStack:
    """A powerful node to load multiple LoRAs with individual controls."""
    
    def __init__(self):
        # Use a single instance for caching benefits
        self.lora_loader = LoraLoader()
    
    @classmethod
    def INPUT_TYPES(cls):
        _strength_inputs = {
            f"strength_{i}": ("FLOAT", {
                "default": 1.0, "min": -2.0, "max": 2.0, "step": 0.01, "forceInput": True,
            }) for i in range(1, 11)
        }
        return {
            "required": {
                "model": ("MODEL",),
                "clip": ("CLIP",),
                "lora_stack": ("STRING", {
                    "multiline": False,
                    "default": "[]",
                    "forceInput": False
                }),
            },
            "optional": _strength_inputs,
        }

    RETURN_TYPES = ("MODEL", "CLIP")
    RETURN_NAMES = ("MODEL", "CLIP") 
    FUNCTION = "load_loras"
    CATEGORY = "loaders"

    def load_loras(self, model, clip, lora_stack="[]", **kwargs):
        """Loads multiple LoRAs based on the configuration."""
        import json
        
        # Debug output
        print(f"\n=== MultiLoRAStack Debug ===")
        print(f"Received lora_stack: {lora_stack}")
        
        # Parse lora_stack if it's a string
        lora_configs = []
        if isinstance(lora_stack, str):
            try:
                lora_configs = json.loads(lora_stack)
                print(f"✅ Parsed {len(lora_configs)} LoRA configs from JSON")
                for i, config in enumerate(lora_configs):
                    print(f"  Config {i}: {config}")
            except Exception as e:
                print(f"❌ Failed to parse JSON: {e}")
                return (model, clip)
        else:
            lora_configs = lora_stack if isinstance(lora_stack, list) else []
        
        applied_count = 0
        current_model = model
        current_clip = clip
        
        # Process each LoRA configuration
        for i, config in enumerate(lora_configs):
            print(f"\n--- Processing LoRA {i+1} ---")
            
            if not isinstance(config, dict):
                print(f"❌ Skipping invalid config at index {i}: {config}")
                continue
                
            # Check if this LoRA should be applied
            if not config.get('on', False):
                print(f"⏸️ Skipping disabled LoRA: {config.get('lora', 'Unknown')}")
                continue
                
            lora_name = config.get('lora')
            if not lora_name or lora_name == "None":
                print(f"❌ Skipping empty/None LoRA")
                continue
                
            strength_model = float(config.get('strength', 1.0))
            # strength_N input overrides the widget value (1-indexed slot position)
            override = kwargs.get(f"strength_{i + 1}")
            if override is not None:
                strength_model = float(override)
            strength_clip = float(config.get('strengthTwo', strength_model))
            
            # Skip if both strengths are 0
            if strength_model == 0 and strength_clip == 0:
                print(f"⏸️ Skipping zero-strength LoRA: {lora_name}")
                continue
            
            print(f"🔄 Attempting to load LoRA:")
            print(f"   Name: {lora_name}")
            print(f"   Model Strength: {strength_model}")
            print(f"   Clip Strength: {strength_clip}")
            
            # Apply the LoRA using ComfyUI's standard loader
            try:
                # Apply the LoRA to current model and clip
                current_model, current_clip = self.lora_loader.load_lora(
                    current_model, current_clip, lora_name, strength_model, strength_clip
                )
                applied_count += 1
                print(f"✅ Successfully applied LoRA: {lora_name}")
                
            except Exception as e:
                print(f"❌ Failed to load LoRA '{lora_name}': {str(e)}")
                print(f"   Error type: {type(e).__name__}")
                import traceback
                print(f"   Full traceback: {traceback.format_exc()}")
        
        print(f"\n=== Summary ===")
        print(f"Applied {applied_count} out of {len(lora_configs)} LoRAs")
        print(f"===================\n")
        
        return (current_model, current_clip)


class MultiLoRAStackModelOnly:
    """A powerful node to load multiple LoRAs with individual controls (Model Only)."""

    def __init__(self):
        # Use a single instance for caching benefits
        self.lora_loader = LoraLoaderModelOnly()

    @classmethod
    def INPUT_TYPES(cls):
        _strength_inputs = {
            f"strength_{i}": ("FLOAT", {
                "default": 1.0, "min": -2.0, "max": 2.0, "step": 0.01, "forceInput": True,
            }) for i in range(1, 11)
        }
        return {
            "required": {
                "model": ("MODEL",),
                "lora_stack": ("STRING", {
                    "multiline": False,
                    "default": "[]",
                    "forceInput": False
                }),
            },
            "optional": _strength_inputs,
        }

    RETURN_TYPES = ("MODEL",)
    RETURN_NAMES = ("MODEL",)
    FUNCTION = "load_loras"
    CATEGORY = "loaders"

    def load_loras(self, model, lora_stack="[]", **kwargs):
        """Loads multiple LoRAs based on the configuration (Model Only)."""
        import json
        
        # Debug output
        print(f"\n=== MultiLoRAStackModelOnly Debug ===")
        print(f"Received lora_stack: {lora_stack}")
        
        # Parse lora_stack if it's a string
        lora_configs = []
        if isinstance(lora_stack, str):
            try:
                lora_configs = json.loads(lora_stack)
                print(f"✅ Parsed {len(lora_configs)} LoRA configs from JSON")
                for i, config in enumerate(lora_configs):
                    print(f"  Config {i}: {config}")
            except Exception as e:
                print(f"❌ Failed to parse JSON: {e}")
                return (model,)
        else:
            lora_configs = lora_stack if isinstance(lora_stack, list) else []
        
        applied_count = 0
        current_model = model
        
        # Process each LoRA configuration
        for i, config in enumerate(lora_configs):
            print(f"\n--- Processing LoRA {i+1} ---")
            
            if not isinstance(config, dict):
                print(f"❌ Skipping invalid config at index {i}: {config}")
                continue
                
            # Check if this LoRA should be applied
            if not config.get('on', False):
                print(f"⏸️ Skipping disabled LoRA: {config.get('lora', 'Unknown')}")
                continue
                
            lora_name = config.get('lora')
            if not lora_name or lora_name == "None":
                print(f"❌ Skipping empty/None LoRA")
                continue
                
            strength_model = float(config.get('strength', 1.0))
            # strength_N input overrides the widget value (1-indexed slot position)
            override = kwargs.get(f"strength_{i + 1}")
            if override is not None:
                strength_model = float(override)

            # Skip if strength is 0
            if strength_model == 0:
                print(f"⏸️ Skipping zero-strength LoRA: {lora_name}")
                continue
            
            print(f"🔄 Attempting to load LoRA (Model Only):")
            print(f"   Name: {lora_name}")
            print(f"   Model Strength: {strength_model}")
            
            # Apply the LoRA using ComfyUI's model-only loader
            try:
                # FIXED: Use the correct method name
                current_model, = self.lora_loader.load_lora_model_only(
                    current_model, lora_name, strength_model
                )
                applied_count += 1
                print(f"✅ Successfully applied LoRA (Model Only): {lora_name}")
                
            except Exception as e:
                print(f"❌ Failed to load LoRA '{lora_name}': {str(e)}")
                print(f"   Error type: {type(e).__name__}")
                import traceback
                print(f"   Full traceback: {traceback.format_exc()}")
        
        print(f"\n=== Summary ===")
        print(f"Applied {applied_count} out of {len(lora_configs)} LoRAs (Model Only)")
        print(f"===================\n")
        
        return (current_model,)


# ---------------------------------------------------------------------------
# Orthogonal stacking helpers
# ---------------------------------------------------------------------------

def _extract_up_down(patch):
    """Return (lora_up, lora_down) from a ComfyUI patch, or (None, None).

    Handles:
      LoRAAdapter object  — comfy.weight_adapter.lora, weights attr holds tensors
      ("lora", (up, down, alpha, dora))  — tagged string format (modern)
      (up, down, ...)                    — tensor-first format (old)
      (alpha_float, (up, down, ...))     — float-prefixed format (some versions)
    """
    # comfy.weight_adapter.lora.LoRAAdapter — newer ComfyUI
    if hasattr(patch, "weights"):
        w = patch.weights
        if isinstance(w, (tuple, list)) and len(w) >= 2:
            return w[0], w[1]
        return None, None

    if isinstance(patch, (list, tuple)) and len(patch) >= 2:
        first = patch[0]

        if isinstance(first, str):
            # ("lora", (up, down, alpha, dora))
            v = patch[1]
            if isinstance(v, (tuple, list)) and len(v) >= 2:
                return v[0], v[1]

        elif isinstance(first, torch.Tensor):
            # (up, down, ...)
            return patch[0], patch[1]

        elif isinstance(first, (int, float)):
            rest = patch[1]
            if isinstance(rest, str) and len(patch) >= 3:
                inner = patch[2]
                if isinstance(inner, (tuple, list)) and len(inner) >= 2:
                    return inner[0], inner[1]
            elif isinstance(rest, (tuple, list)) and len(rest) >= 2:
                if isinstance(rest[0], torch.Tensor):
                    return rest[0], rest[1]
                elif isinstance(rest[0], str) and len(rest) >= 2:
                    inner = rest[1]
                    if isinstance(inner, (tuple, list)) and len(inner) >= 2:
                        return inner[0], inner[1]

    return None, None


def _replace_down(patch, new_down):
    """Return a copy of patch with lora_down replaced by new_down."""
    import copy

    # comfy.weight_adapter.lora.LoRAAdapter
    if hasattr(patch, "weights"):
        new_patch = copy.copy(patch)
        w = patch.weights
        new_patch.weights = (w[0], new_down) + tuple(w[2:])
        return new_patch

    if not isinstance(patch, (list, tuple)) or len(patch) < 2:
        return patch

    first = patch[0]

    if isinstance(first, str):
        v = list(patch[1])
        v[1] = new_down
        return (patch[0], tuple(v))

    elif isinstance(first, torch.Tensor):
        v = list(patch)
        v[1] = new_down
        return tuple(v)

    elif isinstance(first, (int, float)):
        rest = patch[1]
        if isinstance(rest, str) and len(patch) >= 3:
            inner = list(patch[2])
            inner[1] = new_down
            return (patch[0], patch[1], tuple(inner))
        elif isinstance(rest, (tuple, list)) and len(rest) >= 2:
            if isinstance(rest[0], torch.Tensor):
                inner = list(rest)
                inner[1] = new_down
                return (patch[0], tuple(inner))
            elif isinstance(rest[0], str) and len(rest) >= 2:
                inner = list(rest[1])
                inner[1] = new_down
                return (patch[0], (rest[0], tuple(inner)))

    return patch


def _gram_schmidt_ortho_downs(downs_list):
    """Gram-Schmidt orthogonalization of down projection matrices in input space.

    Each successive down matrix has its rows projected onto the orthogonal
    complement of the row spaces already committed by earlier LoRAs, so that
    LoRA_i reads from input directions that LoRA_1..i-1 do not respond to.

    downs_list: list of (rank, in_dim) float tensors
    Returns:    list of orthogonalized tensors, same device/dtype as inputs
    """
    result = []
    Q = None  # (k, in_dim) — accumulated orthonormal row basis

    for down in downs_list:
        orig_dtype = down.dtype
        orig_device = down.device
        d = down.float().to(orig_device)

        if d.dim() != 2:
            result.append(down)
            continue

        if Q is not None and Q.shape[0] > 0:
            # Project rows of d out of the committed subspace
            proj = (d @ Q.T) @ Q          # (r, in_dim)
            d_ortho = d - proj
        else:
            d_ortho = d.clone()

        result.append(d_ortho.to(orig_dtype))

        # Extend Q with the row space of d_ortho
        fro = d_ortho.norm(p="fro")
        if fro.item() < 1e-10:
            continue
        try:
            _, S, Vh = torch.linalg.svd(d_ortho, full_matrices=False)
            # Vh: (r, in_dim) — rows are right singular vectors
            tol = max(S[0].item() * 1e-6, 1e-10) if S.numel() > 0 else 1e-10
            keep = S > tol
            if not keep.any():
                continue
            Q_new = Vh[keep]  # (k', in_dim) orthonormal rows

            if Q is None:
                Q = Q_new
            else:
                # Concatenate and re-orthogonalize
                combined = torch.cat([Q, Q_new], dim=0).T  # (in_dim, k+k')
                Q_upd, R = torch.linalg.qr(combined, mode="reduced")
                r_diag = R.diagonal().abs()
                if r_diag.numel() > 0 and r_diag.max().item() > 0:
                    keep_q = r_diag > r_diag.max().item() * 1e-6
                    Q = Q_upd[:, keep_q].T  # (k_new, in_dim)
                else:
                    Q = Q_upd.T
        except Exception as exc:
            print(f"[MultiLoRAOrtho] SVD/QR failed during basis update: {exc}")

    return result


class MultiLoRAStackModelOnlyOrtho:
    """Multi LoRA Stack with Gram-Schmidt input-space orthogonalization.

    Reduces interference between stacked LoRAs by projecting each LoRA's
    down matrix onto the orthogonal complement of the row spaces already
    claimed by earlier LoRAs. Ports the OrthoLoRA principle from anima_lora
    (Cayley-rotated SVD basis, networks/lora_modules/ortho.py) to
    inference-time stacking without requiring retraining.

    merge_method:
      sequential — identical to the plain MultiLoRAStackModelOnly node.
      ortho_gs   — Gram-Schmidt projection in input space before applying.
    """

    def __init__(self):
        self.lora_loader = LoraLoaderModelOnly()

    @classmethod
    def INPUT_TYPES(cls):
        _strength_inputs = {
            f"strength_{i}": ("FLOAT", {
                "default": 1.0, "min": -2.0, "max": 2.0, "step": 0.01, "forceInput": True,
            }) for i in range(1, 11)
        }
        return {
            "required": {
                "model": ("MODEL",),
                "lora_stack": ("STRING", {
                    "multiline": False,
                    "default": "[]",
                    "forceInput": False,
                }),
                "merge_method": (["ortho_gs", "sequential"], {
                    "default": "ortho_gs",
                }),
            },
            "optional": _strength_inputs,
        }

    RETURN_TYPES = ("MODEL",)
    RETURN_NAMES = ("MODEL",)
    FUNCTION = "load_loras"
    CATEGORY = "loaders"

    def load_loras(self, model, lora_stack="[]", merge_method="ortho_gs", **kwargs):
        import json

        print(f"\n=== MultiLoRAStackModelOnlyOrtho [{merge_method}] ===")

        lora_configs = []
        if isinstance(lora_stack, str):
            try:
                lora_configs = json.loads(lora_stack)
            except Exception as e:
                print(f"❌ JSON parse error: {e}")
                return (model,)
        else:
            lora_configs = lora_stack if isinstance(lora_stack, list) else []

        enabled = []
        for i, cfg in enumerate(lora_configs):
            if not isinstance(cfg, dict) or not cfg.get("on", False):
                continue
            name = cfg.get("lora")
            if not name or name == "None":
                continue
            strength = float(cfg.get("strength", 1.0))
            # strength_N input overrides the widget value (1-indexed slot position)
            override = kwargs.get(f"strength_{i + 1}")
            if override is not None:
                strength = float(override)
            if strength == 0.0:
                continue
            enabled.append((name, strength))

        if not enabled:
            print("No enabled LoRAs.")
            return (model,)

        if merge_method == "sequential" or len(enabled) == 1:
            return self._sequential(model, enabled)

        return self._ortho_gs(model, enabled)

    # ------------------------------------------------------------------
    def _sequential(self, model, lora_list):
        current = model
        for name, strength in lora_list:
            print(f"  sequential: {name} @ {strength}")
            current, = self.lora_loader.load_lora_model_only(current, name, strength)
        return (current,)

    # ------------------------------------------------------------------
    def _ortho_gs(self, model, lora_list):
        import comfy.utils
        import comfy.lora
        import folder_paths

        # Build model key map
        key_map = {}
        try:
            key_map = comfy.lora.model_lora_keys_unet(model.model, key_map)
        except Exception as e:
            print(f"[MultiLoRAOrtho] model_lora_keys_unet failed ({e}); falling back to sequential")
            return self._sequential(model, lora_list)

        # Load all LoRAs into per-key patch dicts
        all_patches = []  # list of ({model_key: patch_tuple}, strength, name)
        for name, strength in lora_list:
            path = folder_paths.get_full_path("loras", name)
            if not path:
                print(f"[MultiLoRAOrtho] ❌ LoRA not found: {name}")
                continue
            try:
                lora_sd = comfy.utils.load_torch_file(path, safe_load=True)
                patches = comfy.lora.load_lora(lora_sd, key_map)
                all_patches.append((patches, strength, name))
                print(f"  loaded: {name} @ {strength} ({len(patches)} keys)")
            except Exception as e:
                print(f"[MultiLoRAOrtho] ❌ Failed to load {name}: {e}")

        if not all_patches:
            return (model,)

        # Group by model key
        key_to_entries = {}  # key -> [(patch_tuple, strength, name)]
        for patches, strength, name in all_patches:
            for key, patch in patches.items():
                key_to_entries.setdefault(key, []).append((patch, strength, name))

        new_model = model.clone()
        ortho_count = 0
        sequential_count = 0

        for key, entries in key_to_entries.items():
            if len(entries) == 1:
                patch, strength, _ = entries[0]
                new_model.add_patches({key: patch}, strength)
                sequential_count += 1
                continue

            # Extract (up, down) from each entry
            components = []
            fallback = False
            for patch, strength, name in entries:
                up, down = _extract_up_down(patch)
                if up is None or down is None or down.dim() != 2:
                    # Non-LoRA or conv patch — skip orthogonalization
                    fallback = True
                    break
                components.append((up, down, patch, strength))

            if fallback:
                for patch, strength, _ in entries:
                    new_model.add_patches({key: patch}, strength)
                sequential_count += len(entries)
                continue

            # Gram-Schmidt on the down matrices
            downs = [c[1] for c in components]
            ortho_downs = _gram_schmidt_ortho_downs(downs)

            for (up, down, patch, strength), ortho_down in zip(components, ortho_downs):
                new_patch = _replace_down(patch, ortho_down.to(down.dtype))
                new_model.add_patches({key: new_patch}, strength)
            ortho_count += 1

        print(f"  ortho keys: {ortho_count}, pass-through: {sequential_count}")
        print("===================\n")
        return (new_model,)


# Register the nodes
NODE_CLASS_MAPPINGS = {
    "MultiLoRAStack": MultiLoRAStack,
    "MultiLoRAStackModelOnly": MultiLoRAStackModelOnly,
    "MultiLoRAStackModelOnlyOrtho": MultiLoRAStackModelOnlyOrtho,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "MultiLoRAStack": "Multi LoRA Stack",
    "MultiLoRAStackModelOnly": "Multi LoRA Stack (Model Only)",
    "MultiLoRAStackModelOnlyOrtho": "Multi LoRA Stack Ortho (Model Only)",
}