from inspect import cleandoc
from comfy_api.latest import io
from .helper_functions import commonLazy
from .modelLikeCommon import calculate_patches


class ModelMathNode(io.ComfyNode):
    """
    This node enables the use of math expressions on Model weights (state_dict).
    Functionally acts as a custom model merge.
    """

    @classmethod
    def define_schema(cls) -> io.Schema:
        return io.Schema(
            node_id="mrmth_ModelMathNode",
            display_name="Model Math",
            category="More math",
            inputs=[
                io.Model.Input(id="a", tooltip="Main model (base)"),
                io.Model.Input(id="b", optional=True, lazy=True, tooltip="Optional 2nd model"),
                io.Model.Input(id="c", optional=True, lazy=True, tooltip="Optional 3rd model"),
                io.Model.Input(id="d", optional=True, lazy=True, tooltip="Optional 4th model"),
                io.Float.Input(id="w", default=0.0, optional=True, lazy=True, force_input=True),
                io.Float.Input(id="x", default=0.0, optional=True, lazy=True, force_input=True),
                io.Float.Input(id="y", default=0.0, optional=True, lazy=True, force_input=True),
                io.Float.Input(id="z", default=0.0, optional=True, lazy=True, force_input=True),
                io.String.Input(id="Model", default="a*(1-w)+b*w", tooltip="Expression to apply on weights"),
                io.Combo.Input(
                    id="length_mismatch",
                    options=["broadcast", "passthrough", "pad"],
                    default="broadcast",
                    tooltip="How to handle mismatched layer counts. For models, this usually defaults to broadcast (zero for missing layers)."
                )
            ],
            outputs=[
                io.Model.Output(),
            ],
        )

    tooltip = cleandoc(__doc__)

    @classmethod
    def check_lazy_status(cls, Model, a, b=[], c=[], d=[], w=0, x=0, y=0, z=0, length_mismatch="broadcast"):
        return commonLazy(Model, a, b, c, d, w, x, y, z)

    @classmethod
    def execute(cls, Model, a, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0.0, length_mismatch="broadcast") -> io.NodeOutput:
        patches = calculate_patches(Model, a, b, c, d, w, x, y, z)
        out_model = a.clone()
        if patches:
            out_model.add_patches(patches, 1.0, 1.0)
        return (out_model,)