#===----------------------------------------------------------------------===//
#
#                                   xara
#                              https://xara.so
#
#===----------------------------------------------------------------------===//
#
# Copyright (c) 2025, OpenSees/Xara Developers
# All rights reserved.  No warranty, explicit or implicit, is provided.
#
# This source code is licensed under the BSD 2-Clause License.
# See LICENSE file or https://opensource.org/licenses/BSD-2-Clause
#===----------------------------------------------------------------------===//
import jax
import inspect
from functools import wraps
from typing import get_args

class ArrayMeta(type):
    def __getitem__(cls, dims):
        if not isinstance(dims, tuple):
            dims = (dims,)
        return type(f"Array{dims}", (cls,), {"__dims__": dims})

class Array(metaclass=ArrayMeta):
    __dims__ = ()



def NodalValue(create_model, output, steps=1):
    node, dof = output

    from xara.para import _Parameter

    pnames = list(inspect.signature(create_model).parameters.keys())

    @jax.custom_vjp
    @wraps(create_model)
    def f(*params):
        m = create_model(*(
            _Parameter(p, index=i, name=n) for (i, p), n in zip(enumerate(params), pnames)
        ))
        m.analyze(steps)
        return m.nodeDisp(node, dof)

    # Forward
    def f_fwd(*params):
        m = create_model(*(
            _Parameter(p, index=i, name=n) for (i, p), n in zip(enumerate(params), pnames)
        ))
        m.analyze(steps)
        y = m.nodeDisp(node, dof)
        return y, m


    # Backward
    def f_bwd(m, cotangent):
        return tuple(
            m.sensNodeDisp(node, dof, tag) * cotangent
            for tag in m._model._openseespy._interp._tcl.call("getParamTags")
        )

    f.defvjp(f_fwd, f_bwd)
    return f


def _extract_parameters(func, params):
    hints = func.__annotations__

    spec = {}
    i = 0
    for j,name in enumerate(inspect.signature(func).parameters.keys()):
        type = hints.get(name, float)
        if type is float:
            spec[name] = {
                "name": name,
                "type": float,
                "index": i
            }
            i += 1

        elif issubclass(type,Array):
            spec[name] = {
                "name": name,
                "type": Array,
                "items": []
            }
            nr, names = type.__dims__
            nr = len(params[j])
            nc = len(names)
            for r in range(nr):
                spec[name]["items"].append([])
                for c in range(nc):
                    spec[name]["items"][-1].append({
                        "name": names[c],
                        "type": float,
                        "index": i
                    })
                    i += 1
        else:
            raise TypeError(f"Unsupported parameter type: {type}")

    return spec

def _shape_parameters(params, pspec):
    from xara.para import _Parameter

    for param, spec in zip(params, pspec.values()):
        if issubclass(spec["type"], Array):
            array = []
            for r in range(len(spec["items"])):
                array.append([])
                for c in range(len(spec["items"][r])):
                    array[-1].append(_Parameter( param[r][c], index=spec["items"][r][c]["index"], name=spec["items"][r][c]["name"]))
            yield array
        elif spec["type"] is float:
            yield _Parameter(param, index=spec["index"], name=spec["name"])
        else:
            raise TypeError(f"Unsupported parameter type: {spec['type']}")


def StaticAnalysis(create_model, output, steps=1):
    node, dof = output


    # @jax.custom_vjp
    @jax.custom_jvp
    @wraps(create_model)
    def f(*params):
        pspec = _extract_parameters(create_model, params)
        m = create_model(*_shape_parameters(params, pspec))
        m._pspec = pspec
        m.analyze(steps)
        return m.nodeDisp(node, dof)

    @f.defjvp
    def f_jvp(primals, tangents):
        # --- Primal pass ---
        pspec = _extract_parameters(create_model, primals)
        m = create_model(*_shape_parameters(primals, pspec))
        m._pspec = pspec
        m.analyze(steps)
        primal_out = m.nodeDisp(node, dof)

        # --- Tangent pass: JVP = Σ (∂u/∂pᵢ) · ṗᵢ ---
        tangent_out = 0.0
        for param, tangent in zip(pspec.values(), tangents):
            if issubclass(param["type"], Array):
                for r, row in enumerate(param["items"]):
                    for c, item in enumerate(row):
                        sens = m.sensNodeDisp(node, dof, item["index"])
                        tangent_out = tangent_out + sens * tangent[r][c]
            elif param["type"] is float:
                sens = m.sensNodeDisp(node, dof, param["index"])
                tangent_out = tangent_out + sens * tangent
            else:
                raise TypeError(f"Unsupported parameter type: {param['type']}")

        return primal_out, tangent_out
    return f

    # Forward
    def f_fwd(*params):
        pspec = _extract_parameters(create_model, params)
        m = create_model(*_shape_parameters(params, pspec))
        m._pspec = pspec
        m.analyze(steps)
        y = m.nodeDisp(node, dof)
        return y, m


    # Backward
    def f_bwd(m, cotangent):
        du = []
        pspec = m._pspec
        for param in pspec.values():
            if issubclass(param["type"], Array):
                array = []
                for r in range(len(param["items"])):
                    array.append([])
                    for c in range(len(param["items"][r])):
                        array[-1].append(m.sensNodeDisp(node, dof, param["items"][r][c]["index"]) * cotangent)
                du.append(array)
            elif param["type"] is float:
                du.append(m.sensNodeDisp(node, dof, param["index"]) * cotangent)
            else:
                raise TypeError(f"Unsupported parameter type: {param['type']}")
        return tuple(du)
        # return tuple(
        #     m.sensNodeDisp(node, dof, tag) * cotangent
        #     for tag in m._model._openseespy._interp._tcl.call("getParamTags")
        # )

    f.defvjp(f_fwd, f_bwd)
    return f

