"""Табличная трикубическая коэнергия: поток и якобиан от одного потенциала.

Узловые W' и dW'/di=(1.5 λd,1.5 λq,λf) приходят из FEMM.
Смешанные производные оцениваются симметрично по соседним узлам.
Это альтернативная интерполяция, а не симметризация чужого якобиана.
"""
from numerics import FluxMap


def basis(t,side,derivative):
    if side==0 and derivative==0:return (2*t**3-3*t*t+1,6*t*t-6*t,12*t-6)
    if side==1 and derivative==0:return (-2*t**3+3*t*t,-6*t*t+6*t,-12*t+6)
    if side==0:return (t**3-2*t*t+t,3*t*t-4*t+1,6*t-4)
    return (t**3-t*t,3*t*t-2*t,6*t-2)


class EnergyFluxMap(FluxMap):
    def __init__(self,axes,values,coenergy):
        super().__init__(axes,values)
        self.coenergy=coenergy;self.weights=(1.5,1.5,1.);self.derivatives={}
        shape=[len(a) for a in axes]
        points=[(i,j,k) for i in range(shape[0]) for j in range(shape[1]) for k in range(shape[2])]
        for index in points:
            i,j,k=index
            self.derivatives[index]={(0,0,0):coenergy[i][j][k],(1,0,0):1.5*values[i][j][k][0],
                                     (0,1,0):1.5*values[i][j][k][1],(0,0,1):values[i][j][k][2]}
        def diff(index,flags,axis):
            lo=list(index);hi=list(index)
            lo[axis]=max(0,index[axis]-1);hi[axis]=min(shape[axis]-1,index[axis]+1)
            return (self.derivatives[tuple(hi)][flags]-self.derivatives[tuple(lo)][flags])/(axes[axis][hi[axis]]-axes[axis][lo[axis]])
        for index in points:
            for a,b in ((0,1),(0,2),(1,2)):
                fa=[0,0,0];fb=[0,0,0];both=[0,0,0];fa[a]=1;fb[b]=1;both[a]=both[b]=1
                self.derivatives[index][tuple(both)]=(diff(index,tuple(fa),b)+diff(index,tuple(fb),a))/2
        for index in points:
            self.derivatives[index][(1,1,1)]=(diff(index,(1,1,0),2)+diff(index,(1,0,1),1)+diff(index,(0,1,1),0))/3

    def potential(self,point):
        # Parent enforces finite inputs and the exact declared domain.
        if len(point)!=3:raise ValueError('Нужны id,iq,If')
        indices=[];widths=[];tables=[]
        for axis,x in zip(self.axes,point):
            if not axis[0]<=x<=axis[-1]:raise ValueError('Экстраполяция карты запрещена')
            k=0
            while k<len(axis)-2 and x>=axis[k+1]:k+=1
            width=axis[k+1]-axis[k];t=(x-axis[k])/width
            indices.append(k);widths.append(width)
            tables.append([[tuple(v*width**derivative/width**order for order,v in enumerate(basis(t,side,derivative)))
                            for derivative in (0,1)] for side in (0,1)])
        value=0.;gradient=[0.,0.,0.];hessian=[[0.]*3 for _ in range(3)]
        for a in (0,1):
            for b in (0,1):
                for c in (0,1):
                    corners=(a,b,c);data=self.derivatives[tuple(indices[j]+corners[j] for j in range(3))]
                    for flags,coefficient in data.items():
                        v=[tables[j][corners[j]][flags[j]] for j in range(3)]
                        value+=coefficient*v[0][0]*v[1][0]*v[2][0]
                        for j in range(3):
                            gradient[j]+=coefficient*v[j][1]*v[(j+1)%3][0]*v[(j+2)%3][0]
                            for k in range(3):
                                term=coefficient
                                for axis in range(3):term*=v[axis][int(axis==j)+int(axis==k)]
                                hessian[j][k]+=term
        return value,gradient,hessian

    def evaluate(self,point):
        _,gradient,hessian=self.potential(point)
        return [gradient[j]/self.weights[j] for j in range(3)],[[hessian[j][k]/self.weights[j] for k in range(3)] for j in range(3)]

    def energy(self,point):
        coenergy,gradient,_=self.potential(point)
        return sum(point[j]*gradient[j] for j in range(3))-coenergy


def from_record(record):
    if record.get('interpolation')=='coenergy-tricubic-hermite':
        return EnergyFluxMap(record['axes_A'],record['values_Wb'],record['coenergy_J'])
    return FluxMap(record['axes_A'],record['values_Wb'])