#!/usr/bin/env python3
"""Independent two-fluid matrices. Units Omega=eta*vK=1, exp(sigma*t+i*k.x)."""
import numpy as np


def equilibrium(stokes,loading,pressure_speed=1):
    denominator=(1+loading)**2+stokes**2
    gas=np.array([2*loading*stokes,-(1+loading+stokes**2),0])*pressure_speed/denominator
    dust=np.array([-2*stokes,-(1+loading),0])*pressure_speed/denominator
    return gas,dust


def incompressible_matrix(stokes,loading,kx,kz,diffusion=0,pressure_speed=1,feedback=True):
    gas,dust=equilibrium(stokes,loading,pressure_speed)
    k=np.array([kx,0.,kz]);norm=np.linalg.norm(k)
    basis=np.array([[kz/norm,0],[0,1],[-kx/norm,0]])
    coriolis=np.array([[0,2,0],[-.5,0,0],[0,0,0]])
    ratio=loading if feedback else 0
    matrix=np.zeros((6,6),dtype=complex)
    matrix[:2,:2]=basis.T@coriolis@basis-(ratio/stokes+1j*kx*gas[0])*np.eye(2)
    matrix[:2,2:5]=ratio/stokes*basis.T
    matrix[:2,5]=ratio/stokes*basis.T@(dust-gas)
    matrix[2:5,:2]=basis/stokes
    matrix[2:5,2:5]=coriolis-(1/stokes+1j*kx*dust[0])*np.eye(3)
    matrix[5,2:5]=-1j*k
    matrix[5,5]=-1j*kx*dust[0]
    # Explicit toy equal diffusivity on all evolved perturbations, not a complete alpha turbulence closure.
    matrix-=diffusion*(kx*kx+kz*kz)*np.eye(6)
    return matrix,basis


def compressible_matrix(stokes,loading,kx,kz,sound_speed=20):
    gas,dust=equilibrium(stokes,loading)
    k=np.array([kx,0.,kz]);coriolis=np.array([[0,2,0],[-.5,0,0],[0,0,0]])
    matrix=np.zeros((8,8),dtype=complex)
    # State gas velocity3, dust velocity3, gas relative density, dust relative density.
    matrix[:3,:3]=coriolis-(loading/stokes+1j*kx*gas[0])*np.eye(3)
    matrix[:3,3:6]=loading/stokes*np.eye(3)
    matrix[:3,6]=-loading/stokes*(dust-gas)-1j*sound_speed**2*k
    matrix[:3,7]=loading/stokes*(dust-gas)
    matrix[3:6,:3]=np.eye(3)/stokes
    matrix[3:6,3:6]=coriolis-(1/stokes+1j*kx*dust[0])*np.eye(3)
    matrix[6,:3]=-1j*k;matrix[6,6]=-1j*kx*gas[0]
    matrix[7,3:6]=-1j*k;matrix[7,7]=-1j*kx*dust[0]
    return matrix


def dominant(matrix):
    values,vectors=np.linalg.eig(matrix)
    index=np.argmax(values.real)
    normalization=vectors[-1,index] if abs(vectors[-1,index])>1e-12 else np.linalg.norm(vectors[:,index])
    vector=vectors[:,index]/normalization
    return values[index],vector,values


if __name__=='__main__':
    for name,stokes,loading,kx,kz,published in [('linA',.1,3,30,30,.4190204),('linB',.1,.2,6,6,.0154764)]:
        matrix,_=incompressible_matrix(stokes,loading,kx,kz)
        value,vector,_=dominant(matrix)
        full,_,_=dominant(compressible_matrix(stokes,loading,kx,kz))
        print(name,'incompressible',value,'compressible',full,'published',published)
        print('vector',vector)


def youdin_goodman_pencil(stokes,loading,kx,kz):
    """Direct COM/relative-velocity equations26-31, with algebraic pressure constraint.
    Separate implementation from the individual-phase transverse projection.
    """
    gas_fraction=1/(1+loading);dust_fraction=1-gas_fraction
    gas,dust=equilibrium(stokes,loading);relative=dust-gas
    k=np.array([kx,0.,kz]);coriolis=np.array([[0,2,0],[-.5,0,0],[0,0,0]])
    def right(state):
        center=state[:3];difference=state[3:6];density=state[6];pressure=state[7]
        f=1j*gas_fraction*((dust_fraction*(k@difference)+gas_fraction*kx*relative[0]*density)*relative+dust_fraction*kx*relative[0]*difference)
        g=1j*kx*relative[0]*((gas_fraction-dust_fraction)*difference-gas_fraction*relative*density)
        center_rhs=coriolis@center-f-1j*k*pressure-density*np.array([2*gas_fraction,0,0])
        difference_rhs=coriolis@difference-1j*kx*relative[0]*center-g-(difference+density*relative)/(gas_fraction*stokes)+1j*k*pressure/gas_fraction
        constraint=k@center-dust_fraction*(k@difference)-gas_fraction*kx*relative[0]*density
        return np.concatenate([center_rhs,difference_rhs,[-1j*(k@center),constraint]])
    operator=np.column_stack([right(np.eye(8,dtype=complex)[:,index])for index in range(8)])
    mass=np.diag([1,1,1,1,1,1,1,0]).astype(complex)
    return operator,mass