from __future__ import print_function, division
import numpy as np
from numpy import require, zeros, array, where, unravel_index, einsum
from pyscf.nao import nao, prod_basis
from pyscf.nao import conv_yzx2xyz_c
from pyscf.dft.rks import _dft_common_init_

#
#
#
class mf(nao):

  def __init__(self, **kw):
    """
    Constructor a mean field class
    store result of a mean-field calc, deliver density matrix etc
    """
    #print(__name__, 'before construct')
    nao.__init__(self, **kw)

    if 'mf' in kw:
      self.init_mo_from_pyscf(**kw)
    elif 'label' in kw: # init KS orbitals with SIESTA
      self.k2xyzw = self.xml_dict["k2xyzw"]
      self.xc_code = 'LDA,PZ' # just a guess...
    elif 'wfsx_fname' in kw: # init KS orbitals with WFSX file from SIESTA output
      self.xc_code = 'LDA,PZ' # just a guess...
    elif 'fireball' in kw: # init KS orbitals with Fireball
      self.init_mo_coeff_fireball(**kw)
      self.xc_code = 'GGA,PBE' # just a guess...
    elif 'gpaw' in kw:
      self.init_mo_coeff_label(**kw)
      self.k2xyzw = np.array([[0.0,0.0,0.0,1.0]])
      self.xc_code = 'LDA,PZ' # just a guess, but in case of GPAW there is a field...
    elif 'openmx' in kw:
      self.xc_code = 'GGA,PBE' # just a guess...
      pass
    else:
      print(__name__, kw.keys())
      raise RuntimeError('unknown constructor')

    #_dft_common_init_(self)

    if self.verbosity>0:
      print(__name__,'\t\t====> self.pseudo: ', self.pseudo)
      print(__name__,'\t\t====> Number of orbitals: ', self.norbs)

    self.init_libnao()

    self.gen_pb = kw['gen_pb'] if 'gen_pb' in kw else True
    if self.gen_pb:
      self.pb = pb = prod_basis(nao=self, **kw)
      self.v_dab = pb.get_dp_vertex_sparse(dtype=self.dtype).tocsr()
      self.cc_da = cc = pb.get_da2cc_sparse(dtype=self.dtype).tocsr()
      self.nprod = self.cc_da.shape[1]
      if self.verbosity>0: print(__name__,'\t\t====> Number of dominant and atom-centered products {}'.format(cc.shape))

      #self.pb.init_prod_basis_pp_batch(nao=self, **kw)

  def init_mo_from_pyscf(self, **kw):
    """ Initializing from a previous pySCF mean-field calc. """
    from pyscf.nao.m_fermi_energy import fermi_energy as comput_fermi_energy
    from pyscf.nao.m_color import color as tc
    self.telec = kw['telec'] if 'telec' in kw else 0.0000317 # 10K
    self.mf = mf = kw['mf']
    self.xc_code = mf.xc if hasattr(mf, 'xc') else 'HF'
    self.k2xyzw = np.array([[0.0,0.0,0.0,1.0]])

    self.mo_energy = np.asarray(mf.mo_energy)
    self.nspin = self.mo_energy.ndim
    assert self.nspin in [1,2]
    nspin,n=self.nspin,self.norbs
    self.mo_energy = require( self.mo_energy.reshape((1, nspin, n)), requirements='CW')
    self.mo_occ = require( mf.mo_occ.reshape((1,nspin,n)), requirements='CW')
    self.mo_coeff =  require(zeros((1,nspin,n,n,1), dtype=self.dtype), requirements='CW')
    conv = conv_yzx2xyz_c(kw['gto'])
    aaux = np.asarray(mf.mo_coeff).reshape((nspin,n,n))
    for s in range(nspin):
      self.mo_coeff[0,s,:,:,0] = conv.conv_yzx2xyz_1d(aaux[s], conv.m_xyz2m_yzx).T

    self.nelec = kw['nelec'] if 'nelec' in kw else np.array([int(s2o.sum()) for s2o in self.mo_occ[0]])
    fermi = comput_fermi_energy(self.mo_energy, sum(self.nelec), self.telec)
    self.fermi_energy = kw['fermi_energy'] if 'fermi_energy' in kw else fermi

  def init_mo_coeff_fireball(self, **kw):
    """ Constructor a mean-field class from the preceeding FIREBALL calculation """
    from pyscf.nao.m_fermi_dirac import fermi_dirac_occupations
    from pyscf.nao.m_fireball_get_eigen_dat import fireball_get_eigen_dat
    from pyscf.nao.m_fireball_hsx import fireball_hsx
    self.telec = kw['telec'] if 'telec' in kw else self.telec
    self.fermi_energy = kw['fermi_energy'] if 'fermi_energy' in kw else self.fermi_energy
    self.mo_energy = require(fireball_get_eigen_dat(self.cd), dtype=self.dtype, requirements='CW')
    ksn2fd = fermi_dirac_occupations(self.telec, self.mo_energy, self.fermi_energy)
    self.mo_occ = (3-self.nspin)*ksn2fd
    if abs(self.nelectron-self.mo_occ.sum())>1e-6: raise RuntimeError("mo_occ wrong?" )
    #print(__name__, ' self.nspecies ', self.nspecies)
    #print(self.sp_mu2j)

    self.hsx = fireball_hsx(self, **kw)
    #print(self.telec)
    #print(self.mo_energy)
    #print(self.fermi_energy)
    #print(__name__, ' sum(self.mo_occ)', sum(self.mo_occ))
    #print(self.mo_occ)

  def diag_check(self, atol=1e-5, rtol=1e-4):
    from pyscf.nao.m_sv_diag import sv_diag
    ksn2e = self.mo_energy
    ac = True
    for k,kvec in enumerate(self.k2xyzw):
      for spin in range(self.nspin):
        e,x = sv_diag(self, kvec=kvec[0:3], spin=spin)
        eref = ksn2e[k,spin,:]
        acks = np.allclose(eref,e,atol=atol,rtol=rtol)
        ac = ac and acks
        if(not acks):
          aerr = sum(abs(eref-e))/len(e)
          print("diag_check: "+bc.RED+str(k)+' '+str(spin)+' '+str(aerr)+bc.ENDC)
    return ac

  def get_occupations(self, telec=None, ksn2e=None, fermi_energy=None):
    """ Compute occupations of electron levels according to Fermi-Dirac distribution """
    from pyscf.nao.m_fermi_dirac import fermi_dirac_occupations
    Telec = self.hsx.telec if telec is None else telec
    ksn2E = self.wfsx.ksn2e if ksn2e is None else ksn2e
    Fermi = self.fermi_energy if fermi_energy is None else fermi_energy
    ksn2fd = fermi_dirac_occupations(Telec, ksn2E, Fermi)
    ksn2fd = (3.0-self.nspin)*ksn2fd
    return ksn2fd

  def init_libnao(self, wfsx=None):
    """ Initialization of data on libnao site """
    from pyscf.nao.m_libnao import libnao
    from pyscf.nao.m_sv_chain_data import sv_chain_data
    from ctypes import POINTER, c_double, c_int64, c_int32, byref

    if wfsx is None:
        data = sv_chain_data(self)
        # (nkpoints, nspin, norbs, norbs, nreim)
        #print(' data ', sum(data))
        size_x = np.array([1, self.nspin, self.norbs, self.norbs, 1], dtype=np.int32)
        libnao.init_sv_libnao_orbs.argtypes = (POINTER(c_double), POINTER(c_int64), POINTER(c_int32))
        libnao.init_sv_libnao_orbs(data.ctypes.data_as(POINTER(c_double)), c_int64(len(data)), size_x.ctypes.data_as(POINTER(c_int32)))
        self.init_sv_libnao = True
    else:
        size_x = np.zeros(len(self.wfsx.x.shape), dtype=np.int32)
        for i, sh in enumerate(self.wfsx.x.shape): size_x[i] = sh

        data = sv_chain_data(self)
        libnao.init_sv_libnao_orbs.argtypes = (POINTER(c_double), POINTER(c_int64), POINTER(c_int32))
        libnao.init_sv_libnao_orbs(data.ctypes.data_as(POINTER(c_double)), c_int64(len(data)), size_x.ctypes.data_as(POINTER(c_int32)))
        self.init_sv_libnao = True

    libnao.init_aos_libnao.argtypes = (POINTER(c_int64), POINTER(c_int64))
    info = c_int64(-999)
    libnao.init_aos_libnao(c_int64(self.norbs), byref(info))
    if info.value!=0: raise RuntimeError("info!=0")
    return self

  def vxc_lil(self, **kw):   # Compute matrix elements of exchange-correlation potential
    from pyscf.nao.m_vxc_lil import vxc_lil
    return vxc_lil(self, deriv=1, **kw)

  def vhartree_pbc(self, dens, **kw):
    """  Compute Hartree potential for the density given in an equidistant grid  """
    from pyscf.nao.m_vhartree_pbc import vhartree_pbc
    return vhartree_pbc(self, dens, **kw)

  def vhartree_pbc_coo(self, density_factors=[1,0], **kw):
    """  Compute matrix elements of Hartree potential for the density given in an equidistant grid  """
    from pyscf.nao.m_vhartree_pbc import vhartree_pbc
    g = self.mesh3d.get_3dgrid()
    f = density_factors
    dens = np.zeros(g.shape)
    if abs(f[0])>0: dens += f[0]*self.dens_elec(g.coords, self.make_rdm1()).reshape(g.shape)
    if abs(f[1])>0: dens += f[1]*self.vna(g.coords,sp2v=self.ao_log.sp2chlocal,sp2rcut=self.ao_log.sp2rcut_chlocal).reshape(g.shape)

    #print(__name__, dens.sum()*self.mesh3d.dv)
    vh = self.vhartree_pbc(dens)
    return self.matelem_int3d_coo(g, vh)

  def dens_elec(self, coords, dm): # Compute electronic density for a given density matrix and on a given set of coordinates
    from pyscf.nao.m_dens_libnao import dens_libnao
    from pyscf.nao.m_init_dm_libnao import init_dm_libnao
    from pyscf.nao.m_init_dens_libnao import init_dens_libnao
    # end of imports
    if not self.init_sv_libnao : raise RuntimeError('not self.init_sv_libnao')
    if init_dm_libnao(dm) is None : raise RuntimeError('init_dm_libnao(dm) is None')
    if init_dens_libnao()!=0 : raise RuntimeError('init_dens_libnao()!=0')
    return dens_libnao(coords, self.nspin)

  def exc(self, dm=None, xc_code=None, **kw):   # Compute exchange-correlation energies
    from pyscf.nao.m_exc import exc
    dm = self.make_rdm1() if dm is None else dm
    xc_code = self.xc_code if xc_code is None else xc_code
    return exc(self, dm, xc_code, **kw)

  def get_init_guess(self, mol=None, key=None):
    """ Compute an initial guess for the density matrix. """
    from pyscf.scf.hf import init_guess_by_minao
    if hasattr(self, 'mol'):
      dm = init_guess_by_minao(self.mol)
    else:
      dm = self.make_rdm1()  # the loaded ks orbitals will be used
      if dm.shape[0:2]==(1,1) and dm.shape[4]==1 : dm = dm.reshape((self.norbs,self.norbs))
    return dm

  def get_hamiltonian(self): # Returns the stored matrix elements of current hamiltonian
    return self.hsx.spin2h4_csr

  def dos(self, comegas, **kw):
    """ Ordinary Density of States (from the current mean-field eigenvalues) """
    from pyscf.nao.scf_dos import scf_dos
    return scf_dos(self, comegas, **kw)

  def pdos(self, comegas, **kw):
    """
    Partial Density of States (resolved in angular momentum of atomic orbitals)
    """
    import pyscf.nao.m_dos_pdos_ldos as mpdos
    return mpdos.pdos(self, comegas, **kw)

  def lsoa_dos(self, comegas, **kw):
    """
    Partial Density of States (contributions from a given list of atoms)
    """
    import pyscf.nao.m_dos_pdos_ldos as mpdos
    return mpdos.lsoa_dos(self, comegas, **kw)

  def gdos(self, comegas, **kw):
    """
    Some molecular orbital population analysis
    """
    import pyscf.nao.m_dos_pdos_ldos as mpdos
    return mpdos.gdos(self, comegas, **kw)

  def read_wfsx(self, fname, **kw):
    """
    An occasional reading of the SIESTA's .WFSX file
    """
    from pyscf.nao.m_siesta_wfsx import siesta_wfsx_c
    from pyscf.nao.m_siesta2blanko_denvec import _siesta2blanko_denvec
    from pyscf.nao.m_fermi_dirac import fermi_dirac_occupations

    self.wfsx = siesta_wfsx_c(fname=fname, **kw)

    assert self.nkpoints == self.wfsx.nkpoints
    assert self.norbs == self.wfsx.norbs
    assert self.nspin == self.wfsx.nspin
    orb2m = self.get_orb2m()
    for k in range(self.nkpoints):
      for s in range(self.nspin):
        for n in range(self.norbs):
          _siesta2blanko_denvec(orb2m, self.wfsx.x[k,s,n,:,:])

    self.mo_coeff = require(self.wfsx.x, dtype=self.dtype, requirements='CW')
    self.mo_energy = require(self.wfsx.ksn2e, dtype=self.dtype, requirements='CW')
    self.telec = kw['telec'] if 'telec' in kw else self.hsx.telec
    self.nelec = kw['nelec'] if 'nelec' in kw else self.hsx.nelec
    self.fermi_energy = kw['fermi_energy'] if 'fermi_energy' in kw else self.fermi_energy
    ksn2fd = fermi_dirac_occupations(self.telec, self.mo_energy, self.fermi_energy)
    self.mo_occ = (3-self.nspin)*ksn2fd
    return self

  def plot_contour(self, w=0.0):
    """
      Plot contour with poles of Green's function in the self-energy
      SelfEnergy(w) = G(w+w')W(w')
      with respect to w' = Re(w')+Im(w')
      Poles of G(w+w') are located: w+w'-(E_n-Fermi)+i*eps sign(E_n-Fermi)==0 ==>
      w'= (E_n-Fermi) - w -i eps sign(E_n-Fermi)
    """
    try :
      import matplotlib.pyplot as plt
      from matplotlib.patches import Arc, Arrow
    except:
      print('no matplotlib?')
      return

    fig,ax = plt.subplots()
    fe = self.fermi_energy
    ee = self.mo_energy
    iee = 0.5-np.array(ee>fe)
    eew = ee-fe-w
    ax.plot(eew, iee, 'r.', ms=10.0)
    pp = list()
    pp.append(Arc((0,0),4,4,angle=0, linewidth=2, theta1=0, theta2=90, zorder=2, color='b'))
    pp.append(Arc((0,0),4,4,angle=0, linewidth=2, theta1=180, theta2=270, zorder=2, color='b'))
    pp.append(Arrow(0,2,0,-4,width=0.2, color='b', hatch='o'))
    pp.append(Arrow(-2,0,4,0,width=0.2, color='b', hatch='o'))
    for p in pp: ax.add_patch(p)
    ax.set_aspect('equal')
    ax.grid(True, which='both')
    ax.axhline(y=0, color='k')
    ax.axvline(x=0, color='k')
    plt.ylim(-3.0,3.0)
    plt.show()

  def get_vertex_pov(self):
    """ Computes the occupied-virtual-product basis vertex"""
    assert hasattr(self, 'pb')
    pab = self.pb.get_ac_vertex_array()
    pov = list()
    nprd = pab.shape[0]
    for s,occ in enumerate(self.mo_occ):
      no = np.count_nonzero(occ>0.0)
      nv = np.count_nonzero(occ==0.0)
      assert nv+no==self.mo_coeff.shape[-2]
      pov.append(np.zeros([nprd,no,nv], dtype=self.dtype))
      pov[s] = np.einsum('oa,pab,vb->pov', self.mo_coeff[0,s,0:no,:,0], pab, self.mo_coeff[0,s,no:,:,0])
    return pov

  def nonin_osc_strength(self):
    from scipy.sparse import spmatrix
    """ Computes the non-interacting oscillator strengths and energies """

    x,y,z = map(spmatrix.toarray, self.dipole_coo())
    i2d = array((x,y,z))
    n = self.mo_occ.shape[-1]

    p = zeros((len(comega)), dtype=np.complex128) # result to accumulate

    for s in range(self.nspin):
      o,e,cc = self.mo_occ[0,s],self.mo_energy[0,s],self.mo_coeff[0,s,:,:,0]
      oo1,ee1 = np.subtract.outer(o,o).reshape(n*n), np.subtract.outer(e,e).reshape(n*n)
      idx = unravel_index( np.intersect1d(where(oo1<0.0), where(ee1<eemax)), (n,n))
      ivrt,iocc = array(list(set(idx[0]))), array(list(set(idx[1])))
      voi2d = einsum('nia,ma->nmi', einsum('iab,nb->nia', i2d, cc[ivrt]), cc[iocc])
      t2osc = 2.0/3.0*einsum('voi,voi->vo', voi2d, voi2d)
      t2w =  np.subtract.outer(e[ivrt],e[iocc])
      t2o = -np.subtract.outer(o[ivrt],o[iocc])

      for iw,w in enumerate(comega):
        p[iw] += 0.5*(t2osc*((t2o/(w-t2w))-(t2o/(w+t2w)))).sum()
    return p

  def polariz_nonin_ave_matelem(self, comega):
    from scipy.sparse import spmatrix
    """ Computes the non-interacting optical polarizability via the dipole matrix elements."""

    x,y,z = map(spmatrix.toarray, self.dipole_coo())
    i2d = array((x,y,z))
    n = self.mo_occ.shape[-1]
    eemax = max(comega.real)+20.0*max(comega.imag)

    p = zeros((len(comega)), dtype=np.complex128) # result to accumulate

    #print(__name__, 'Fermi energy', self.fermi_energy)
    #np.set_printoptions(linewidth=1000)
    for s in range(self.nspin):
      o,e,cc = self.mo_occ[0,s],self.mo_energy[0,s],self.mo_coeff[0,s,:,:,0]
      #print(o[:10])
      #print(e[:10])

      oo1,ee1 = np.subtract.outer(o,o).reshape(n*n), np.subtract.outer(e,e).reshape(n*n)
      idx = unravel_index( np.intersect1d(where(oo1<0.0), where(ee1<eemax)), (n,n))
      ivrt,iocc = array(list(set(idx[0]))), array(list(set(idx[1])))
      voi2d = einsum('nia,ma->nmi', einsum('iab,nb->nia', i2d, cc[ivrt]), cc[iocc])
      t2osc = 2.0/3.0*einsum('voi,voi->vo', voi2d, voi2d)
      t2w =  np.subtract.outer(e[ivrt],e[iocc])
      t2o = -np.subtract.outer(o[ivrt],o[iocc])

      for iw,w in enumerate(comega):
        p[iw] += 0.5*(t2osc*((t2o/(w-t2w))-(t2o/(w+t2w)))).sum()

    return p

  def spin_square(self, mo_coeff=None, mo_occ=None):
    from functools import reduce

    mo_coeff = self.mo_coeff if mo_coeff is None else mo_coeff
    mo_occ = self.mo_occ if mo_occ is None else mo_occ

    if self.nspin==1:
      mo_a = mo_coeff[0,0,mo_occ[0,0]>0,:,0]
      mo_b = mo_coeff[0,0,mo_occ[0,0]>1,:,0]
    elif self.nspin==2:
      mo_a = mo_coeff[0,0,mo_occ[0,0]>0,:,0]
      mo_b = mo_coeff[0,1,mo_occ[0,1]>0,:,0]

    nocc_a, nocc_b = mo_a.shape[0], mo_b.shape[0]
    over = self.overlap_coo().toarray()
    s = reduce(np.dot, (mo_a, over, mo_b.T))
    ssxy = (nocc_a+nocc_b) * 0.5 - (s.conj()*s).sum()
    ssz = (nocc_b-nocc_a)**2 * 0.25
    ss = (ssxy + ssz).real
    s = np.sqrt(ss+0.25) - 0.5

    return ss, s*2+1

#
# Example of reading pySCF mean-field calculation.
#
if __name__=="__main__":
  from pyscf import gto, scf as scf_gto
  from pyscf.nao import nao, mf
  """ Interpreting small Gaussian calculation """
  mol = gto.M(atom='O 0 0 0; H 0 0 1; H 0 1 0; Be 1 0 0', basis='ccpvdz') # coordinates in Angstrom!
  dft = scf_gto.RKS(mol)
  dft.kernel()

  sv = mf(mf=dft, gto=dft.mol, rcut_tol=1e-9, nr=512, rmin=1e-6)

  print(sv.ao_log.sp2norbs)
  print(sv.ao_log.sp2nmult)
  print(sv.ao_log.sp2rcut)
  print(sv.ao_log.sp_mu2rcut)
  print(sv.ao_log.nr)
  print(sv.ao_log.rr[0:4], sv.ao_log.rr[-1:-5:-1])
  print(sv.ao_log.psi_log[0].shape, sv.ao_log.psi_log_rl[0].shape)
  print(dir(sv.pb))
  print(sv.pb.norbs)
  print(sv.pb.npdp)
  print(sv.pb.c2s[-1])


import itertools

import numpy as np
import torch as th
from ttools.training import ModelInterface

from . import utils

class ReconstructionInterface(ModelInterface):
    def __init__(self, model, args, vertex_idxs, face_idxs, junctions,
                 edge_data, vertex_t, adjacencies, junction_order,
                 template_normals, symmetries=None):
        self.model = model

        self.vertex_idxs = vertex_idxs
        self.face_idxs = face_idxs
        self.junctions = junctions
        self.edge_data = edge_data
        self.vertex_t = vertex_t
        self.junction_order = junction_order
        self.template_normals = template_normals[None]

        self.args = args

        self.n_samples_per_loop_side = int(
            np.ceil(np.sqrt(args.n_samples / face_idxs.shape[0])))

        self.optimizer = th.optim.Adam(self.model.parameters(), lr=args.lr)

        self.edge_idxs = [[0, 1, 2, 3], [3, 4, 5, 6],
                          [6, 7, 8, 9], [9, 10, 11, 0]]

        self.nonadjacent_patch_pairs = []
        self.adjacent_patch_pairs = []
        for i1, i2 in list(
                itertools.combinations(range(face_idxs.shape[0]), 2)):
            if args.wheels and (
                    i1 in [5, 11, 17, 23] or i2 in [5, 11, 17, 23]):
                continue
            if i2 in adjacencies[i1]:
                e1 = adjacencies[i1][i2]
                self.adjacent_patch_pairs.append((i1, i2, e1))
            else:
                self.nonadjacent_patch_pairs.append((i1, i2))

        self.d_points_to_tris = utils.PointToTriangleDistance.apply

        p_edge0, p_edge1, p_edge2, p_edge3 = [], [], [], []
        for i in range(self.n_samples_per_loop_side):
            for j in range(self.n_samples_per_loop_side):
                n = self.n_samples_per_loop_side

                if i > 0:
                    p_edge0.append(i + j*self.n_samples_per_loop_side)
                if i < n-1:
                    p_edge2.append(i + j*self.n_samples_per_loop_side)
                if j > 0:
                    p_edge3.append(i + j*self.n_samples_per_loop_side)
                if j < n-1:
                    p_edge1.append(i + j*self.n_samples_per_loop_side)
        self.grid_point_edges = (p_edge0, p_edge1, p_edge2, p_edge3)

        self.triangulation = []
        t_edge0, t_edge1, t_edge2, t_edge3 = [], [], [], []
        for i in range(self.n_samples_per_loop_side-1):
            for j in range(self.n_samples_per_loop_side-1):
                n = self.n_samples_per_loop_side-1

                if i > 1:
                    t_edge0.extend(
                        [len(self.triangulation), len(self.triangulation)+1])
                if i < n-2:
                    t_edge2.extend(
                        [len(self.triangulation), len(self.triangulation)+1])
                if j > 1:
                    t_edge3.extend(
                        [len(self.triangulation), len(self.triangulation)+1])
                if j < n-2:
                    t_edge1.extend(
                        [len(self.triangulation), len(self.triangulation)+1])

                self.triangulation.extend(
                    [[i + j*self.n_samples_per_loop_side,
                      i + (j+1)*self.n_samples_per_loop_side,
                      i + (j+1)*self.n_samples_per_loop_side + 1],
                     [i + j*self.n_samples_per_loop_side,
                      i + j*self.n_samples_per_loop_side + 1,
                      i + (j+1)*self.n_samples_per_loop_side + 1]])
        self.triangulation_edges = (t_edge0, t_edge1, t_edge2, t_edge3)

        _loss = Loss(args, self.triangulation_edges, self.triangulation,
                     self.grid_point_edges, self.n_samples_per_loop_side,
                     self.d_points_to_tris, self.edge_idxs,
                     self.nonadjacent_patch_pairs, self.adjacent_patch_pairs,
                     self.template_normals, symmetries)
        self._compute_losses = th.nn.DataParallel(
            _loss) if args.cuda else _loss

        if args.cuda:
            self.model.cuda()
            self._compute_losses.cuda()

    def forward(self, batch):
        ims = batch['ims']
        if self.args.cuda:
            ims = ims.cuda()

        params = self.model(ims)
        vertices, patches = utils.process_patches(
            params, self.vertex_idxs, self.face_idxs, self.edge_data,
            self.junctions, self.junction_order, self.vertex_t)

        st = th.empty(patches.shape[0], patches.shape[1],
                      self.n_samples_per_loop_side**2, 2).uniform_().to(params)

        points = utils.coons_sample(st[..., 0], st[..., 1], patches)
        normals = utils.coons_normals(st[..., 0], st[..., 1], patches)
        mtds = utils.coons_mtds(st[..., 0], st[..., 1], patches)

        return {'patches': patches, 'points': points, 'normals': normals,
                'mtds': mtds, 'st': st, 'params': params, 'vertices': vertices}

    def training_step(self, batch):
        self.model.train()
        self.optimizer.zero_grad()

        losses_dict = self._compute_losses(batch, self.forward(batch))
        loss = losses_dict['loss']
        loss.mean().backward()
        self.optimizer.step()

        return {k: v.mean().item() for k, v in losses_dict.items()}

    def init_validation(self):
        losses = ['loss', 'chamferloss', 'normalsloss', 'collisionloss',
                  'planarloss', 'templatenormalsloss', 'symmetryloss']
        ret = {loss: 0 for loss in losses}
        ret['count'] = 0
        return ret

    def validation_step(self, batch, running_data):
        self.model.eval()
        count = running_data['count']
        n = batch['ims'].shape[0]
        losses_dict = self._compute_losses(batch, self.forward(batch))
        loss = losses_dict['loss']
        chamferloss = losses_dict['chamferloss']
        normalsloss = losses_dict['normalsloss']
        collisionloss = losses_dict['collisionloss']
        planarloss = losses_dict['planarloss']
        templatenormalsloss = losses_dict['templatenormalsloss']
        symmetryloss = losses_dict['symmetryloss']
        return {
            'loss': (running_data['loss']*count +
                     loss.mean().item()*n) / (count+n),
            'chamferloss': (running_data['chamferloss']*count +
                            chamferloss.mean().item()*n) / (count+n),
            'normalsloss': (running_data['normalsloss']*count +
                            normalsloss.mean().item()*n) / (count+n),
            'collisionloss': (running_data['collisionloss']*count +
                              collisionloss.mean().item()*n) / (count+n),
            'planarloss': (running_data['planarloss']*count +
                           planarloss.mean().item()*n) / (count+n),
            'templatenormalsloss': (running_data['templatenormalsloss']*count +
                                    templatenormalsloss.mean().item()*n)
            / (count+n),
            'symmetryloss': (running_data['symmetryloss']*count
                             + symmetryloss.mean().item()*n) / (count+n),
            'count': count + n
        }

wheel_idxs = list(range(24))
no_wheel_idxs = list(range(24, 43))

class Loss(th.nn.Module):
    def __init__(self, args, triangulation_edges, triangulation,
                 grid_point_edges, n_samples_per_loop_side, d_points_to_tris,
                 edge_idxs, nonadjacent_patch_pairs, adjacent_patch_pairs,
                 template_normals, symmetries):
        super(Loss, self).__init__()
        self.args = args
        self.triangulation_edges = triangulation_edges
        self.triangulation = triangulation
        self.grid_point_edges = grid_point_edges
        self.d_points_to_tris = d_points_to_tris
        self.edge_idxs = edge_idxs
        self.nonadjacent_patch_pairs = nonadjacent_patch_pairs
        self.adjacent_patch_pairs = adjacent_patch_pairs
        self.symmetries = symmetries

        linspace = th.linspace(0, 1, n_samples_per_loop_side)
        s_grid, t_grid = th.meshgrid(linspace, linspace)
        self.s_grid = th.nn.Parameter(s_grid.flatten(), requires_grad=False)
        self.t_grid = th.nn.Parameter(t_grid.flatten(), requires_grad=False)
        self.template_normals = th.nn.Parameter(
            template_normals, requires_grad=False)

    def forward(self, batch, fwd_data):
        patches = fwd_data['patches']  # [b, n_patches, 12, 3]
        points = fwd_data['points']
        normals = fwd_data['normals']
        mtds = fwd_data['mtds']
        vertices = fwd_data['vertices']

        target_points = batch['points'].to(points)  # [b, n_points, 3]
        target_normals = batch['normals'].to(normals)
        if self.args.wheels:
            wheel_target_points = batch['wheel_points'].to(points)
            wheel_target_normals = batch['wheel_normals'].to(normals)

        b, n_patches, _, _ = patches.shape

        st = fwd_data['st']

        if self.symmetries is not None:
            xs, ys = self.symmetries
            xs_ = vertices[:, xs]
            ys_ = vertices[:, ys] * vertices.new_tensor([[[-1, 1, 1]]])
            symmetryloss = th.sum((xs_ - ys_)**2, dim=-1).mean()
        else:
            symmetryloss = patches.new_zeros(1)

        if self.args.wheels:
            mtds = mtds.view(b, -1)

            wheel_points = points[:, wheel_idxs]
            no_wheel_points = points[:, no_wheel_idxs]

            wheel_normals = normals[:, wheel_idxs]
            no_wheel_normals = normals[:, no_wheel_idxs]

            wheel_chamferloss_a, wheel_chamferloss_b, wheel_normalsloss_a, \
                wheel_normalsloss_b = utils.compute_chamfer_losses(
                    wheel_points, wheel_normals, wheel_target_points,
                    wheel_target_normals, self.args.w_normals > 0)
            no_wheel_chamferloss_a, no_wheel_chamferloss_b, \
                no_wheel_normalsloss_a, no_wheel_normalsloss_b = \
                utils.compute_chamfer_losses(no_wheel_points, no_wheel_normals,
                                             target_points, target_normals,
                                             self.args.w_normals > 0)

            chamferloss_a = th.cat(
                [wheel_chamferloss_a, no_wheel_chamferloss_a], dim=1)
            normalsloss_a = th.cat(
                [wheel_normalsloss_a, no_wheel_normalsloss_a], dim=1)

            ratio = batch['ratio']
            n_wheel_pts = wheel_target_points.shape[1]
            n_no_wheel_pts = target_points.shape[1]
            n_pts = n_wheel_pts + n_no_wheel_pts
            ratio_wheel = ratio * n_pts/n_wheel_pts
            ratio_no_wheel = (1-ratio) * n_pts/n_no_wheel_pts

            chamferloss_a = th.sum(mtds*chamferloss_a, dim=-1) / mtds.sum(-1)
            chamferloss_b = ratio_wheel * \
                wheel_chamferloss_b.mean(
                    1) + ratio_no_wheel * no_wheel_chamferloss_b.mean(1)
            chamferloss = ((chamferloss_a+chamferloss_b).mean() / 2).view(1)

            normalsloss_a = th.sum(mtds*normalsloss_a, dim=-1) / mtds.sum(-1)
            normalsloss_b = ratio_wheel * \
                wheel_normalsloss_b.mean(
                    1) + ratio_no_wheel * no_wheel_normalsloss_b.mean(1)
            normalsloss = ((normalsloss_a+normalsloss_b).mean() / 2).view(1)

        elif self.args.seperate_turbines:
            turbine_idxs = [x for x in range(76) if x not in [72, 73, 74, 75]]
            no_turbine_idxs = [x for x in range(76)
                               if x not in [10, 12, 14, 17, 19, 20, 21, 24, 25,
                                            27, 28, 29, 30, 31, 32, 33, 34,
                                            66]]

            turbine_points = points[batch['turbines']]
            turbine_points = turbine_points[:, turbine_idxs]
            no_turbine_points = points[batch['turbines']]
            no_turbine_points = no_turbine_points[:, no_turbine_idxs]

            turbine_normals = normals[batch['turbines']]
            turbine_normals = turbine_normals[:, turbine_idxs]
            no_turbine_normals = normals[batch['turbines']]
            no_turbine_normals = no_turbine_normals[:, no_turbine_idxs]

            turbine_mtds = mtds[batch['turbines']]
            turbine_mtds = turbine_mtds[:, turbine_idxs]
            no_turbine_mtds = mtds[batch['turbines']]
            no_turbine_mtds = no_turbine_mtds[:, no_turbine_idxs]

            turbine_patches = patches[batch['turbines']]
            turbine_patches = turbine_patches[:, turbine_idxs]
            no_turbine_patches = patches[batch['turbines']]
            no_turbine_patches = no_turbine_patches[:, no_turbine_idxs]

            turbine_target_points = target_points[batch['turbines']]
            no_turbine_target_points = target_points[batch['turbines']]

            turbine_target_normals = target_normals[batch['turbines']]
            no_turbine_target_normals = target_normals[batch['turbines']]

            if turbine_points.shape[0] > 0:
                turbine_mtds = turbine_mtds.view(turbine_mtds.shape[0], -1)
                turbine_chamferloss_a, turbine_chamferloss_b, \
                    turbine_normalsloss_a, turbine_normalsloss_b = \
                    utils.compute_chamfer_losses(turbine_points,
                                                 turbine_normals,
                                                 turbine_target_points,
                                                 turbine_target_normals,
                                                 self.args.w_normals > 0)
                turbine_chamferloss_a = th.sum(
                    turbine_mtds*turbine_chamferloss_a,
                    dim=-1) / turbine_mtds.sum(-1)
                turbine_chamferloss_b = turbine_chamferloss_b.mean(1)
                turbine_chamferloss = (
                    (turbine_chamferloss_a +
                     turbine_chamferloss_b).mean() / 2).view(1)
                turbine_normalsloss_a = th.sum(
                    turbine_mtds*turbine_normalsloss_a,
                    dim=-1) / turbine_mtds.sum(-1)
                turbine_normalsloss_b = turbine_normalsloss_b.mean(-1)
                turbine_normalsloss = (
                    (turbine_normalsloss_a +
                     turbine_normalsloss_b).mean() / 2).view(1)
            else:
                turbine_chamferloss = turbine_normalsloss = points.new_zeros(1)

            if no_turbine_points.shape[0] > 0:
                no_turbine_mtds = no_turbine_mtds.view(
                    no_turbine_mtds.shape[0], -1)
                no_turbine_chamferloss_a, no_turbine_chamferloss_b, \
                    no_turbine_normalsloss_a, no_turbine_normalsloss_b = \
                    utils.compute_chamfer_losses(no_turbine_points,
                                                 no_turbine_normals,
                                                 no_turbine_target_points,
                                                 no_turbine_target_normals,
                                                 self.args.w_normals > 0)
                no_turbine_chamferloss_a = th.sum(
                    no_turbine_mtds*no_turbine_chamferloss_a,
                    dim=-1) / no_turbine_mtds.sum(-1)
                no_turbine_chamferloss_b = no_turbine_chamferloss_b.mean(1)
                no_turbine_chamferloss = (
                    (no_turbine_chamferloss_a +
                     no_turbine_chamferloss_b).mean() / 2).view(1)
                no_turbine_normalsloss_a = th.sum(
                    no_turbine_mtds*no_turbine_normalsloss_a,
                    dim=-1) / no_turbine_mtds.sum(-1)
                no_turbine_normalsloss_b = no_turbine_normalsloss_b.mean(-1)
                no_turbine_normalsloss = (
                    (no_turbine_normalsloss_a +
                     no_turbine_normalsloss_b).mean() / 2).view(1)
            else:
                no_turbine_chamferloss = no_turbine_normalsloss = \
                    points.new_zeros(1)

            ratio_turbine = turbine_points.shape[0] / b
            chamferloss = ratio_turbine*turbine_chamferloss + \
                (1-ratio_turbine)*no_turbine_chamferloss
            normalsloss = ratio_turbine*turbine_normalsloss + \
                (1-ratio_turbine)*no_turbine_normalsloss

            del turbine_normals, turbine_target_normals, no_turbine_normals, \
                no_turbine_target_normals, turbine_points, \
                turbine_target_points, no_turbine_points, \
                no_turbine_target_points
        else:
            mtds = mtds.view(b, -1)

            chamferloss_a, chamferloss_b, normalsloss_a, normalsloss_b = \
                utils.compute_chamfer_losses(
                    points, normals, target_points, target_normals,
                    self.args.w_normals > 0)
            chamferloss_a = th.sum(mtds*chamferloss_a, dim=-1) / mtds.sum(-1)
            chamferloss_b = chamferloss_b.mean(1)
            chamferloss = ((chamferloss_a+chamferloss_b).mean() / 2).view(1)
            normalsloss_a = th.sum(mtds*normalsloss_a, dim=-1) / mtds.sum(-1)
            normalsloss_b = normalsloss_b.mean(-1)
            normalsloss = ((normalsloss_a+normalsloss_b).mean() / 2).view(1)

        mtds = mtds.view(b, n_patches, -1)

        if self.args.w_templatenormals:
            templatenormalsloss = th.sum(
                (self.template_normals - normals)**2, dim=-1)
            templatenormalsloss = th.sum(
                mtds*templatenormalsloss, dim=-1) / mtds.sum(-1)
        else:
            templatenormalsloss = th.zeros_like(chamferloss)

        del target_normals, normals, target_points

        if self.args.w_planar > 0:
            planarloss = utils.planar_patch_loss(st, points, mtds)
        else:
            planarloss = th.zeros_like(chamferloss)

        del points, mtds

        if self.args.w_collision > 0:
            collisionloss = chamferloss.new_zeros([b, 0])
            grid_points = utils.coons_sample(self.s_grid, self.t_grid, patches)
            triangles = grid_points[:, :, self.triangulation]

            i1s, i2s, e1s = zip(*self.adjacent_patch_pairs)
            points1 = grid_points[:, i1s]
            point_idxs = th.tensor([self.grid_point_edges[e]
                                    for e in e1s]).to(points1.device)
            point_idxs = point_idxs[None, :, :, None].expand(b, -1, -1, 3)
            points1 = th.gather(points1, 2, point_idxs)
            points2 = grid_points[:, i2s]

            triangles1 = triangles[:, i1s]
            triangle_idxs = th.tensor(
                [self.triangulation_edges[e] for e in e1s]
            ).to(triangles1.device)
            triangle_idxs = triangle_idxs[None, :, :,
                                          None, None].expand(b, -1, -1, 3, 3)
            triangles1 = th.gather(triangles1, 2, triangle_idxs)
            triangles2 = triangles[:, i2s]

            idxs = utils.bboxes_intersect(
                points1, points2, dim=2).any(0).nonzero().squeeze(1)
            n_adjacent_intersections = idxs.shape[0]

            if n_adjacent_intersections > 0:
                points1 = points1[:, idxs].view([-1] + list(points1.shape[2:]))
                points2 = points2[:, idxs].view([-1] + list(points2.shape[2:]))
                triangles1 = triangles1[:, idxs].view(
                    [-1] + list(triangles1.shape[2:]))
                triangles2 = triangles2[:, idxs].view(
                    [-1] + list(triangles2.shape[2:]))
                d1 = self.d_points_to_tris(points1, triangles2)
                d2 = self.d_points_to_tris(points2, triangles1)
                d = th.min(d1, d2).view(b, -1)
                collisionloss = th.cat(
                    [collisionloss, th.exp(-(d/self.args.sigma_collision)**2)],
                    dim=1)

            i1s, i2s = zip(*self.nonadjacent_patch_pairs)
            idxs = utils.bboxes_intersect(
                grid_points[:, i1s], grid_points[:, i2s], dim=2
            ).any(0).nonzero().squeeze(1)
            n_nonadjacent_intersections = idxs.shape[0]
            i1s = th.tensor(i1s).to(grid_points.device)[idxs]
            i2s = th.tensor(i2s).to(grid_points.device)[idxs]

            if n_nonadjacent_intersections > 0:
                points1 = grid_points[:, i1s].view(
                    [-1] + list(grid_points.shape[2:]))
                points2 = grid_points[:, i2s].view(
                    [-1] + list(grid_points.shape[2:]))
                triangles1 = triangles[:, i1s].view(
                    [-1] + list(triangles.shape[2:]))
                triangles2 = triangles[:, i2s].view(
                    [-1] + list(triangles.shape[2:]))
                d1 = self.d_points_to_tris(points1, triangles2)
                d2 = self.d_points_to_tris(points2, triangles1)
                d = th.min(d1, d2).view(b, -1)
                collisionloss = th.cat(
                    [collisionloss, th.exp(-(d/self.args.sigma_collision)**2)],
                    dim=1)

            del triangles

            if n_adjacent_intersections + n_nonadjacent_intersections > 0:
                collisionloss = collisionloss.sum(-1).mean()
            else:
                collisionloss = th.zeros_like(chamferloss)
        else:
            collisionloss = th.zeros_like(chamferloss)

        loss = chamferloss + self.args.w_normals*normalsloss + \
            self.args.w_collision*collisionloss + \
            self.args.w_planar*planarloss + \
            self.args.w_templatenormals*templatenormalsloss + \
            self.args.w_symmetry*symmetryloss

        return {
            'loss': loss,
            'chamferloss': chamferloss,
            'normalsloss': normalsloss,
            'collisionloss': collisionloss,
            'planarloss': planarloss,
            'templatenormalsloss': templatenormalsloss,
            'symmetryloss': symmetryloss,
        }


# Copyright Contributors to the Pyro-Cov project.
# SPDX-License-Identifier: Apache-2.0

import datetime
import functools
import logging
import math
import pickle
import re
import warnings
from collections import Counter, OrderedDict, defaultdict
from timeit import default_timer
from typing import List

import numpy as np
import pyro
import pyro.distributions as dist
import torch
from pyro import poutine
from pyro.infer import SVI, JitTrace_ELBO, Trace_ELBO
from pyro.infer.autoguide import (
    AutoDelta,
    AutoGuideList,
    AutoLowRankMultivariateNormal,
    AutoNormal,
    AutoStructured,
)
from pyro.infer.reparam import LocScaleReparam
from pyro.ops.streaming import CountMeanVarianceStats, StatsOfDict
from pyro.optim import ClippedAdam
from pyro.poutine.util import site_is_subsample

import pyrocov.geo

from . import pangolin, sarscov2
from .util import pearson_correlation

logger = logging.getLogger(__name__)

# Reasonable values might be week (7), fortnight (14), or month (28)
TIMESTEP = 14  # in days
GENERATION_TIME = 5.5  # in days
START_DATE = "2019-12-01"

def date_range(stop):
    start = datetime.datetime.strptime(START_DATE, "%Y-%m-%d")
    step = datetime.timedelta(days=TIMESTEP)
    return np.array([start + step * t for t in range(stop)])

def get_fine_regions(columns, min_samples=50):
    """
    Select regions that have at least ``min_samples`` samples.
    Remaining regions will be coarsely aggregated up to country level.
    """
    # Count number of samples in each subregion.
    counts = Counter()
    for location in columns["location"]:
        parts = location.split("/")
        if len(parts) < 2:
            continue
        parts = tuple(p.strip() for p in parts[:3])
        counts[parts] += 1

    # Select fine countries.
    return frozenset(parts for parts, count in counts.items() if count >= min_samples)

def rank_loo_lineages(
    full_dataset: dict,
    full_result: dict,
    min_samples: int = 100,
) -> List[str]:
    """
    Compute a list of lineages ranked in descending order of how much their
    growth rate differs from their parents' growth rate. This is used in growth
    rate leave-one-out prediction experiments.
    """
    # Decompress lineage names before computing parents.
    lineage_id_inv = [
        pangolin.decompress(name) for name in full_dataset["lineage_id_inv"]
    ]
    lineage_id = {name: i for i, name in enumerate(lineage_id_inv)}
    ancestors = set(lineage_id)

    # Filter to often-observed lineages.
    weekly_strains = full_dataset["weekly_strains"]  # [T, P, S]
    lineage_counts = weekly_strains.sum([0, 1])  # [S]
    lineages = []
    for c, child in enumerate(lineage_id_inv):
        if child in ("A", "B", "B.1"):
            continue  # ignore very early lineages
        if lineage_counts[c] < min_samples:
            continue  # ignore rare lineages
        lineages.append(child)

    # Sort leaf nodes by their distance from parent.
    rate_loc = full_result["median"]["rate_loc"]
    ranked_lineages = []
    for child in lineages:
        # Allow grandparent to adopt orphan, since e.g. B.1.617 is missing from
        # lineage_id, but B.1.617.2 is very important.
        parent = pangolin.get_most_recent_ancestor(child, ancestors)
        assert parent is not None
        c = lineage_id[child]
        p = lineage_id[parent]
        gap = (rate_loc[c] - rate_loc[p]).abs().item()
        ranked_lineages.append((gap, child))
    ranked_lineages.sort(reverse=True)

    # Compress lineage names before returning.
    return [pangolin.compress(name) for gap, name in ranked_lineages]

def load_gisaid_data(
    *,
    device="cpu",
    include={},
    exclude={},
    end_day=None,
    gisaid_columns_filename="results/gisaid.columns.pkl",
    nextclade_features_filename="results/nextclade.features.pt",
) -> dict:
    """
    Loads the two files gisaid_columns_filename and nextclade_features_filename,
    converts teh input to PyTorch tensors and truncates the data according to
    ``include`` and ``exclude``.

    Keyword arguments:
    device -- torch device to use
    include --
    exclude --
    end_day -- last day to include
    gisaid_columns_filename --
    nextclade_features_filename --
    """
    logger.info("Loading data")
    include = include.copy()
    exclude = exclude.copy()

    if end_day:
        logger.info(f"Load gisaid data end_day: {end_day}")

    # Load ``gisaid_columns_filename``
    with open(gisaid_columns_filename, "rb") as f:
        columns = pickle.load(f)

    logger.info("Training on {} rows with columns:".format(len(columns["day"])))
    logger.info(", ".join(columns.keys()))

    # Filter regions to at least 50 sample and aggregate rest to country level
    fine_regions = get_fine_regions(columns)

    # Filter features into numbers of mutations and possibly genes.
    aa_features = torch.load(nextclade_features_filename)
    mutations = aa_features["mutations"]
    features = aa_features["features"].to(
        device=device, dtype=torch.get_default_dtype()
    )
    keep = [m.count(",") == 0 for m in mutations]  # restrict to single mutations
    if include.get("gene"):
        re_gene = re.compile(include.pop("gene"))
        keep = [k and bool(re_gene.search(m)) for k, m in zip(keep, mutations)]
    if exclude.get("gene"):
        re_gene = re.compile(exclude.pop("gene"))
        keep = [k and not re_gene.search(m) for k, m in zip(keep, mutations)]
    if include.get("region"):
        gene, region = include.pop("region")
        lb, ub = sarscov2.GENE_STRUCTURE[gene][region]
        for i, m in enumerate(mutations):
            g, m = m.split(":")
            if g != gene:
                keep[i] = False
                continue
            match = re.search("[0-9]+", m)
            assert match is not None
            pos = int(match.group())
            if not (lb < pos <= ub):
                keep[i] = False
    mutations = [m for k, m in zip(keep, mutations) if k]
    if mutations:
        features = features[:, keep]
    else:
        warnings.warn("No mutations selected; using empty features")
        mutations = ["S:D614G"]  # bogus
        features = features[:, :1] * 0
    logger.info("Loaded {} feature matrix".format(" x ".join(map(str, features.shape))))

    # Aggregate regions

    # Get lineages
    lineages = list(map(pangolin.compress, columns["lineage"]))
    lineage_id_inv = list(map(pangolin.compress, aa_features["lineages"]))
    lineage_id = {k: i for i, k in enumerate(lineage_id_inv)}

    sparse_data: dict = Counter()
    location_id: dict = OrderedDict()

    # Set of lineages that are skipped
    skipped = set()

    # Generate sparse_data
    for virus_name, day, location, lineage in zip(
        columns["virus_name"], columns["day"], columns["location"], lineages
    ):
        if lineage not in lineage_id:
            if lineage not in skipped:
                skipped.add(lineage)
                logger.warning(f"WARNING skipping unsampled lineage {lineage}")
            continue

        # Filter by include/exclude
        row = {
            "virus_name": virus_name,
            "location": location,
            "day": day,
            "lineage": pangolin.compress(lineage),
        }
        if not all(re.search(v, row[k]) for k, v in include.items()):
            continue
        if any(re.search(v, row[k]) for k, v in exclude.items()):
            continue

        # Filter by day
        if end_day is not None:
            if day > end_day:
                continue

        # preprocess parts
        parts = location.split("/")
        if len(parts) < 2:
            continue
        parts = tuple(p.strip() for p in parts[:3])
        if len(parts) == 3 and parts not in fine_regions:
            parts = parts[:2]
        location = " / ".join(parts)

        p = location_id.setdefault(location, len(location_id))
        s = lineage_id[lineage]
        t = day // TIMESTEP
        sparse_data[t, p, s] += 1

    # Generate weekly_strains tensor from sparse_data
    if end_day is not None:
        T = 1 + end_day // TIMESTEP
    else:
        T = 1 + max(columns["day"]) // TIMESTEP

    P = len(location_id)
    S = len(lineage_id)
    weekly_strains = torch.zeros(T, P, S)
    for (t, p, s), n in sparse_data.items():
        weekly_strains[t, p, s] = n

    logger.info(f"Dataset size [T x P x S] {T} x {P} x {S}")

    logger.info(
        f"Keeping {int(weekly_strains.sum())}/{len(lineages)} rows "
        f"(dropped {len(lineages) - int(weekly_strains.sum())})"
    )

    # Filter regions.
    num_times_observed = (weekly_strains > 0).max(2).values.sum(0)
    ok_regions = (num_times_observed >= 2).nonzero(as_tuple=True)[0]
    ok_region_set = set(ok_regions.tolist())
    logger.info(f"Keeping {len(ok_regions)}/{weekly_strains.size(1)} regions")
    weekly_strains = weekly_strains.index_select(1, ok_regions)
    locations = [k for k, v in location_id.items() if v in ok_region_set]
    location_id = OrderedDict(zip(locations, range(len(locations))))

    # Construct region-local time scales centered around observations.
    num_obs = weekly_strains.sum(-1)
    local_time = torch.arange(float(len(num_obs))) * TIMESTEP / GENERATION_TIME
    local_time = local_time[:, None]
    local_time = local_time - (local_time * num_obs).sum(0) / num_obs.sum(0)

    return {
        "location_id": location_id,
        "mutations": mutations,
        "weekly_strains": weekly_strains,
        "features": features,
        "lineage_id": lineage_id,
        "lineage_id_inv": lineage_id_inv,
        "local_time": local_time,
    }

def subset_gisaid_data(
    gisaid_dataset: dict,
    location_queries=None,
    max_strains=math.inf,
) -> dict:
    """
    Selects a small subset of data for exploratory fitting of a small model.
    This is not used in the final published results.
    """
    old = gisaid_dataset
    new = old.copy()

    # Select locations.
    if location_queries is not None:
        locations = sorted(
            {
                location
                for location in new["location_id"]
                if any(q in location for q in location_queries)
            }
        )
        ids = torch.tensor([old["location_id"][location] for location in locations])
        new["location_id"] = {name: i for i, name in enumerate(locations)}
        new["weekly_strains"] = new["weekly_strains"].index_select(1, ids)
        new["local_time"] = new["local_time"].index_select(1, ids)

    # Select strains.
    if new["weekly_strains"].size(-1) > max_strains:
        ids = (
            new["weekly_strains"]
            .sum([0, 1])
            .sort(0, descending=True)
            .indices[:max_strains]
        )
        new["weekly_strains"] = new["weekly_strains"].index_select(-1, ids)
        new["features"] = new["features"].index_select(0, ids)
        new["lineage_id_inv"] = [new["lineage_id_inv"][i] for i in ids.tolist()]
        new["lineage_id"] = {name: i for i, name in enumerate(new["lineage_id_inv"])}

    # Select mutations.
    gaps = new["features"].max(0).values - new["features"].min(0).values
    ids = (gaps >= 0.5).nonzero(as_tuple=True)[0]
    new["mutations"] = [new["mutations"][i] for i in ids.tolist()]
    new["features"] = new["features"].index_select(-1, ids)

    logger.info(
        "Selected {}/{} places, {}/{} strains, {}/{} mutations, {}/{} samples".format(
            len(new["location_id"]),
            len(old["location_id"]),
            len(new["lineage_id"]),
            len(old["lineage_id"]),
            len(new["mutations"]),
            len(old["mutations"]),
            int(new["weekly_strains"].sum()),
            int(old["weekly_strains"].sum()),
        )
    )

    return new

def load_jhu_data(gisaid_data: dict) -> dict:
    """
    Load case count time series.

    This is used for plotting but is not used for fitting a model.
    """
    # Load raw JHU case count data.
    us_cases_df = pyrocov.geo.read_csv("time_series_covid19_confirmed_US.csv")
    global_cases_df = pyrocov.geo.read_csv("time_series_covid19_confirmed_global.csv")
    daily_cases = torch.cat(
        [
            pyrocov.geo.pd_to_torch(us_cases_df, columns=slice(11, None)),
            pyrocov.geo.pd_to_torch(global_cases_df, columns=slice(4, None)),
        ]
    ).T
    logger.info(
        "Loaded {} x {} daily case data, totaling {}".format(
            *daily_cases.shape, daily_cases[-1].sum().item()
        )
    )

    # Convert JHU locations to GISAID locations.
    locations = list(gisaid_data["location_id"])
    matrix = pyrocov.geo.gisaid_to_jhu_location(locations, us_cases_df, global_cases_df)
    assert matrix.shape == (len(locations), daily_cases.shape[-1])
    daily_cases = daily_cases @ matrix.T
    daily_cases[1:] -= daily_cases[:-1].clone()  # cumulative -> density
    daily_cases.clamp_(min=0)
    assert daily_cases.shape[1] == len(gisaid_data["location_id"])

    # Convert daily counts to TIMESTEP counts (e.g. weekly).
    start_date = datetime.datetime.strptime(START_DATE, "%Y-%m-%d")
    jhu_start_date = pyrocov.geo.parse_date(us_cases_df.columns[11])
    assert start_date < jhu_start_date
    dt = (jhu_start_date - start_date).days
    T = len(gisaid_data["weekly_strains"])
    weekly_cases = daily_cases.new_zeros(T, len(locations))
    for w in range(TIMESTEP):
        t0 = (w + dt) // TIMESTEP
        source = daily_cases[w::TIMESTEP]
        destin = weekly_cases[t0 : t0 + len(source)]
        destin[:] += source[: len(destin)]
    assert weekly_cases.sum() > 0

    return {
        "daily_cases": daily_cases.clamp(min=0),
        "weekly_cases": weekly_cases.clamp(min=0),
    }

def model(dataset, model_type, *, forecast_steps=None):
    """
    Bayesian regression model of lineage portions as a function of mutation features.

    This function can be run in two different modes:
    - During training, ``forecast_steps=None`` and the model is conditioned on
      observed data.
    - During prediction (after training), the likelihood statement is omitted
      and instead a ``probs`` tensor is recorded; this is the predicted lineage
      portions in each (time, regin) bin.
    """
    # Tensor shapes are commented at at the end of some lines.
    features = dataset["features"]
    local_time = dataset["local_time"][..., None]  # [T, P, 1]
    T, P, _ = local_time.shape
    S, F = features.shape
    if forecast_steps is None:  # During inference.
        weekly_strains = dataset["weekly_strains"]
        assert weekly_strains.shape == (T, P, S)
    else:  # During prediction.
        T = T + forecast_steps
        t0 = local_time[0]
        dt = local_time[1] - local_time[0]
        local_time = t0 + dt * torch.arange(float(T))[:, None, None]
        assert local_time.shape == (T, P, 1)
    strain_plate = pyro.plate("strain", S, dim=-1)
    place_plate = pyro.plate("place", P, dim=-2)
    time_plate = pyro.plate("time", T, dim=-3)

    # Configure reparametrization (which does not affect model density).
    reparam = {}
    if "reparam" in model_type:
        local_time = local_time + pyro.param(
            "local_time", lambda: torch.zeros(P, S)
        )  # [T, P, S]
        reparam["coef"] = LocScaleReparam()
        if "skip" not in model_type:
            reparam["rate_loc"] = LocScaleReparam()
        reparam["init_loc"] = LocScaleReparam()
        reparam["rate"] = LocScaleReparam()
        reparam["init"] = LocScaleReparam()
    with poutine.reparam(config=reparam):

        # Sample global random variables.
        coef_scale = pyro.sample("coef_scale", dist.LogNormal(-4, 2))
        if "skip" not in model_type:
            rate_loc_scale = pyro.sample("rate_loc_scale", dist.LogNormal(-4, 2))
        init_loc_scale = pyro.sample("init_loc_scale", dist.LogNormal(0, 2))
        rate_scale = pyro.sample("rate_scale", dist.LogNormal(-4, 2))
        init_scale = pyro.sample("init_scale", dist.LogNormal(0, 2))
        if "poisson" in model_type:
            pois_loc = pyro.sample("pois_loc", dist.Normal(0, 4))
            pois_scale = pyro.sample("pois_scale", dist.LogNormal(0, 4))

        # Assume relative growth rate depends strongly on mutations and weakly
        # on strain and place. Assume initial infections depend strongly on
        # strain and place.
        Dist = dist.Logistic if "sparse" in model_type else dist.Normal
        coef = pyro.sample("coef", Dist(torch.zeros(F), coef_scale).to_event(1))  # [F]
        with strain_plate:
            rate_loc_loc = 0.01 * coef @ features.T
            if "skip" in model_type:
                rate_loc = pyro.deterministic("rate_loc", rate_loc_loc)  # [S]
            else:
                rate_loc = pyro.sample(
                    "rate_loc", dist.Normal(rate_loc_loc, rate_loc_scale)
                )  # [S]
            init_loc = pyro.sample("init_loc", dist.Normal(0, init_loc_scale))  # [S]
        with place_plate, strain_plate:
            rate = pyro.sample("rate", dist.Normal(rate_loc, rate_scale))  # [P, S]
            init = pyro.sample("init", dist.Normal(init_loc, init_scale))  # [P, S]

        # Finally observe counts.
        logits = init + rate * local_time  # [T, P, S]
        if forecast_steps is None:  # During inference.
            if "poisson" in model_type:
                with time_plate, place_plate:
                    pois = pyro.sample("pois", dist.LogNormal(pois_loc, pois_scale))
                # This softmax() breaks the strain_plate, but is more
                # numerically stable than exp(). AutoGaussian inference will be
                # approximate with softmax(), but would be intractable with
                # exp() and a second_strain_plate.
                lambda_ = (pois * logits.softmax(-1)).clamp_(min=1e-6)
                with time_plate, place_plate, strain_plate:
                    pyro.sample(
                        "obs", dist.Poisson(lambda_), obs=weekly_strains
                    )  # [T, P, S]
            else:
                with time_plate, place_plate:
                    pyro.sample(
                        "obs",
                        dist.Multinomial(
                            logits=logits[..., None, :], validate_args=False
                        ),
                        obs=weekly_strains[..., None, :],
                    )  # [T, P, 1, S]
        else:  # During prediction.
            with time_plate, place_plate, strain_plate:
                pyro.deterministic("probs", logits.softmax(-1))

class InitLocFn:
    """
    Initializer for latent variables.

    This is passed as the ``init_loc_fn`` to guides.
    """

    def __init__(self, dataset):
        # Initialize init.
        init = dataset["weekly_strains"].sum(0)  # [P, S]
        init.add_(1 / init.size(-1)).div_(init.sum(-1, True))
        init.log_().sub_(init.median(-1, True).values)
        self.init = init  # [P, S]
        self.init_decentered = init / 2
        self.init_loc = init.mean(0)  # [S]
        self.init_loc_decentered = self.init_loc / 2
        assert not torch.isnan(self.init).any()
        self.pois = dataset["weekly_strains"].sum(-1, True).clamp(min=0.1)  # [T, P, 1]
        logger.info(f"init stddev = {self.init.std():0.3g}")

    def __call__(self, site):
        name = site["name"]
        shape = site["fn"].shape()
        if hasattr(self, name):
            result = getattr(self, name)
            assert result.shape == shape
            return result
        if name in ("coef_scale", "init_scale", "init_loc_scale"):
            return torch.ones(shape)
        if name == "logits_scale":
            return torch.full(shape, 0.002)
        if name in ("rate_loc_scale", "rate_scale", "place_scale", "strain_scale"):
            return torch.full(shape, 0.01)
        if name in (
            "rate_loc",
            "rate_loc_decentered",
            "coef",
            "coef_decentered",
            "rate",
            "rate_decentered",
        ):
            return torch.rand(shape).sub_(0.5).mul_(0.01)
        if name == "coef_loc":
            return torch.rand(shape).sub_(0.5).mul_(0.01).add_(1.0)
        if name == "pois_loc":
            return self.pois.log().mean()
        if name == "pois_scale":
            return self.pois.log().std()
        if name == "pois":
            return self.pois
        raise ValueError(f"InitLocFn found unhandled site {repr(name)}; please update.")

class Guide(AutoGuideList):
    """
    Custom guide for large-scale inference.

    This combines a low-rank multivariate normal guide over small variables
    with a mean field guide over remaining latent variables.
    """

    def __init__(self, model, init_loc_fn, init_scale, rank):
        super().__init__(model)

        # Jointly estimate globals, mutation coefficients, and strain coefficients.
        mvn = [
            "coef_scale",
            "rate_loc_scale",
            "init_loc_scale",
            "rate_scale",
            "init_scale",
            "coef",
            "coef_decentered",
            "rate_loc",
            "rate_loc_decentered",
            "init_loc",
            "init_loc_decentered",
        ]
        self.append(
            AutoLowRankMultivariateNormal(
                poutine.block(model, expose=mvn),
                init_loc_fn=init_loc_fn,
                init_scale=init_scale,
                rank=rank,
            )
        )
        model = poutine.block(model, hide=mvn)

        # Mean-field estimate all remaining latent variables.
        self.append(AutoNormal(model, init_loc_fn=init_loc_fn, init_scale=init_scale))

@torch.no_grad()
@poutine.mask(mask=False)
def predict(
    model,
    guide,
    dataset,
    model_type,
    *,
    num_samples=1000,
    vectorize=None,
    save_params=("rate", "init", "probs"),
    forecast_steps=0,
) -> dict:
    def get_conditionals(data):
        trace = poutine.trace(poutine.condition(model, data)).get_trace(
            dataset, model_type, forecast_steps=forecast_steps
        )
        return {
            name: site["value"].detach()
            for name, site in trace.nodes.items()
            if site["type"] == "sample" and not site_is_subsample(site)
            if name != "obs"
        }

    # Compute median point estimate.
    result: dict = defaultdict(dict)
    for name, value in get_conditionals(guide.median(dataset)).items():
        if value.numel() < 1e5 or name in save_params:
            result["median"][name] = value

    # Compute moments.
    save_params = {
        k for k, v in result["median"].items() if v.numel() < 1e5 or k in save_params
    }
    if vectorize is None:
        vectorize = result["median"]["probs"].numel() < 1e6
    if vectorize:
        with pyro.plate("particles", num_samples, dim=-4):
            samples = get_conditionals(guide())
        for k, v in samples.items():
            if k in save_params:
                result["mean"][k] = v.mean(0).squeeze()
                result["std"][k] = v.std(0).squeeze()
    else:
        stats = StatsOfDict({k: CountMeanVarianceStats for k in save_params})
        for _ in range(num_samples):
            stats.update(get_conditionals(guide()))
            print(".", end="", flush=True)
        for name, stats_ in stats.get().items():
            if "mean" in stats_:
                result["mean"][name] = stats_["mean"]
            if "variance" in stats_:
                result["std"][name] = stats_["variance"].sqrt()
    return dict(result)

def fit_svi(
    dataset: dict,
    *,
    model_type: str,
    guide_type: str,
    cond_data={},
    forecast_steps=0,
    learning_rate=0.05,
    learning_rate_decay=0.1,
    num_steps=3001,
    num_samples=1000,
    clip_norm=10.0,
    rank=200,
    jit=True,
    log_every=50,
    seed=20210319,
    check_loss=False,
) -> dict:
    """
    Fits a variational posterior using stochastic variational inference (SVI).
    """
    start_time = default_timer()

    logger.info(f"Fitting {guide_type} guide via SVI")
    pyro.set_rng_seed(seed)
    pyro.clear_param_store()
    param_store = pyro.get_param_store()

    # Initialize guide so we can count parameters and register hooks.
    cond_data = {k: torch.as_tensor(v) for k, v in cond_data.items()}
    model_ = poutine.condition(model, cond_data)
    init_loc_fn = InitLocFn(dataset)
    if guide_type == "map":
        guide = AutoDelta(model_, init_loc_fn=init_loc_fn)
    elif guide_type == "normal":
        guide = AutoNormal(model_, init_loc_fn=init_loc_fn, init_scale=0.01)
    elif guide_type == "full":
        guide = AutoLowRankMultivariateNormal(
            model_, init_loc_fn=init_loc_fn, init_scale=0.01, rank=rank
        )
    elif guide_type == "structured":
        guide = AutoStructured(
            model_,
            init_loc_fn=init_loc_fn,
            init_scale=0.01,
            conditionals=defaultdict(
                lambda: "normal",
                rate_scale="delta",
                init_loc_scale="delta",
                init_scale="delta",
                coef="mvn",
                coef_decentered="mvn",
            ),
        )
    elif guide_type == "gaussian":
        from pyro.infer.autoguide import AutoGaussian

        guide = AutoGaussian(
            model_, init_loc_fn=init_loc_fn, init_scale=0.01, backend="funsor"
        )
    else:
        guide = Guide(model_, init_loc_fn=init_loc_fn, init_scale=0.01, rank=rank)
    # This initializes the guide:
    latent_shapes = {k: v.shape for k, v in guide(dataset, model_type).items()}
    latent_numel = {k: v.numel() for k, v in latent_shapes.items()}
    logger.info(
        "\n".join(
            [f"Model has {sum(latent_numel.values())} latent variables of shapes:"]
            + [f" {k} {tuple(v)}" for k, v in latent_shapes.items()]
        )
    )
    param_shapes = {k: v.shape for k, v in pyro.get_param_store().named_parameters()}
    param_numel = {k: v.numel() for k, v in param_shapes.items()}
    logger.info(
        "\n".join(
            [f"Guide has {sum(param_numel.values())} parameters of shapes:"]
            + [f" {k} {tuple(v)}" for k, v in param_shapes.items()]
        )
    )

    # Log gradient norms during inference.
    series: dict = defaultdict(list)

    def hook(g, series):
        series.append(torch.linalg.norm(g.reshape(-1), math.inf).item())

    for name, value in pyro.get_param_store().named_parameters():
        value.register_hook(functools.partial(hook, series=series[name]))

    def optim_config(param_name):
        config: dict = {
            "lr": learning_rate,
            "lrd": learning_rate_decay ** (1 / num_steps),
            "clip_norm": clip_norm,
        }
        scalars = [k for k, v in latent_numel.items() if v == 1]
        if any("locs." + s in name for s in scalars):
            config["lr"] *= 0.2
        elif "scales" in param_name:
            config["lr"] *= 0.1
        elif "scale_tril" in param_name:
            config["lr"] *= 0.05
        elif "factors" in param_name:
            config["lr"] *= 0.05
        elif "weight_" in param_name:
            config["lr"] *= 0.01
        elif "weight" in param_name:
            config["lr"] *= 0.03
        elif "_centered" in param_name:
            config["lr"] *= 0.1
        return config

    optim = ClippedAdam(optim_config)
    Elbo = JitTrace_ELBO if jit else Trace_ELBO
    elbo = Elbo(max_plate_nesting=3, ignore_jit_warnings=True)
    svi = SVI(model_, guide, optim, elbo)
    losses = []
    num_obs = dataset["weekly_strains"].count_nonzero()
    for step in range(num_steps):
        loss = svi.step(dataset=dataset, model_type=model_type)
        assert not math.isnan(loss)
        losses.append(loss)
        median = guide.median()
        for name, value in median.items():
            if value.numel() == 1:
                series[name].append(float(value))
        if log_every and step % log_every == 0:
            logger.info(
                " ".join(
                    [f"step {step: >4d} L={loss / num_obs:0.6g}"]
                    + [
                        "{}={:0.3g}".format(
                            "".join(p[0] for p in k.split("_")).upper(), v.item()
                        )
                        for k, v in median.items()
                        if v.numel() == 1
                    ]
                )
            )
        if check_loss and step >= 50:
            prev = torch.tensor(losses[-50:-25], device="cpu").median().item()
            curr = torch.tensor(losses[-25:], device="cpu").median().item()
            assert (curr - prev) < num_obs, "loss is increasing"

    result = predict(
        model_,
        guide,
        dataset,
        model_type,
        num_samples=num_samples,
        forecast_steps=forecast_steps,
    )
    result["losses"] = losses
    series["loss"] = losses
    result["series"] = dict(series)
    result["params"] = {
        k: v.detach().float().cpu().clone()
        for k, v in param_store.items()
        if v.numel() < 1e7
    }
    result["walltime"] = default_timer() - start_time
    return result

@torch.no_grad()
def log_stats(dataset: dict, result: dict) -> dict:
    """
    Logs statistics of predictions and model fit in the ``result`` of
    ``fit_svi()``.

    :param dict dataset: The dataset dictionary.
    :param dict result: The output of :func:`fit_svi`.
    :returns: A dictionary of statistics.
    """
    stats = {k: float(v) for k, v in result["median"].items() if v.numel() == 1}
    stats["loss"] = float(np.median(result["losses"][-100:]))
    mutations = dataset["mutations"]
    mean = result["mean"]["coef"].cpu()
    if not mean.shape:
        return stats  # Work around error in map estimation.

    # Statistical significance.
    std = result["std"]["coef"].cpu()
    sig = mean.abs() / std
    logger.info(f"|μ|/σ [median,max] = [{sig.median():0.3g},{sig.max():0.3g}]")
    stats["|μ|/σ median"] = sig.median()
    stats["|μ|/σ max"] = sig.max()

    # Effects of individual mutations.
    for name in ["S:D614G", "S:N501Y", "S:E484K", "S:L452R"]:
        if name not in mutations:
            continue
        i = mutations.index(name)
        m = mean[i] * 0.01
        s = std[i] * 0.01
        logger.info(f"ΔlogR({name}) = {m:0.3g} ± {s:0.2f}")
        stats[f"ΔlogR({name}) mean"] = m
        stats[f"ΔlogR({name}) std"] = s

    # Growth rates of individual lineages.
    try:
        i = dataset["lineage_id"]["A"]
        rate_A = result["mean"]["rate"][..., i].mean(0)
    except KeyError:
        rate_A = result["mean"]["rate"].median()
    for s in ["B.1.1.7", "B.1.617.2"]:
        i = dataset["lineage_id"][s]
        rate = result["median"]["rate"][..., i].mean()
        R_RA = (rate - rate_A).exp()
        logger.info(f"R({s})/R(A) = {R_RA:0.3g}")
        stats[f"R({s})/R(A)"] = R_RA

    # Accuracy of mutation-only model, ie without region-local effects.
    true = dataset["weekly_strains"] + 1e-20  # avoid nans
    counts = true.sum(-1, True)
    true_probs = true / counts
    local_time = dataset["local_time"][..., None]
    if "local_time" in result["params"]:
        local_time = local_time + result["params"]["local_time"].to(local_time.device)
    rate = 0.01 * result["median"]["coef"] @ dataset["features"].T
    pred = result["median"]["init"] + rate * local_time
    pred -= pred.logsumexp(-1, True)  # apply log sigmoid function
    kl = true.mul(true_probs.log() - pred).sum(-1)
    kl = stats["naive KL"] = kl.sum() / counts.sum()  # in units of nats / observation
    error = (pred.exp() - true_probs) * counts ** 0.5  # scaled by Poisson stddev
    mae = stats["naive MAE"] = error.abs().sum(-1).mean()
    rmse = stats["naive RMSE"] = error.square().sum(-1).mean().sqrt()
    logger.info(f"naive KL = {kl:0.4g}, MAE = {mae:0.4g}, RMSE = {rmse:0.4g}")

    # Posterior predictive error.
    pred = result["median"]["probs"][: len(true)] + 1e-20  # truncate, avoid nans
    kl = true.mul(true_probs.log() - pred.log()).sum([0, -1])
    error = (pred - true_probs) * counts ** 0.5  # scaled by Poisson stddev
    mae = error.abs().mean(0)  # average over time
    mse = error.square().mean(0)  # average over time
    stats["MAE"] = mae.sum(-1).mean()  # average over region
    stats["RMSE"] = mse.sum(-1).mean().sqrt()  # root average over region
    stats["KL"] = kl.sum() / counts.sum()  # in units of nats / observation
    logger.info("KL = {KL:0.4g}, MAE = {MAE:0.4g}, RMSE = {RMSE:0.4g}".format(**stats))

    # Examine the MSE and RMSE over a few regions of interest.
    queries = {
        "England": ["B.1.1.7"],
        # "England": ["B.1.1.7", "B.1.177", "B.1.1", "B.1"],
        # "USA / California": ["B.1.1.7", "B.1.429", "B.1.427", "B.1.2", "B.1", "P.1"],
    }
    for place, strains in queries.items():
        matches = [p for name, p in dataset["location_id"].items() if place in name]
        if not matches:
            continue
        assert len(matches) == 1, matches
        p = matches[0]
        stats[f"{place} KL"] = kl[p].sum() / true[:, p].sum()
        stats[f"{place} MAE"] = mae[p].sum()
        stats[f"{place} RMSE"] = mse[p].sum().sqrt()
        logger.info(
            "{}\tKL = {:0.3g}, MAE = {:0.3g}, RMSE = {:0.3g}".format(
                place,
                stats[f"{place} KL"],
                stats[f"{place} MAE"],
                stats[f"{place} RMSE"],
            )
        )

        for strain in strains:
            s = dataset["lineage_id"][strain]
            stats[f"{place} {strain} MAE"] = mae[p, s]
            stats[f"{place} {strain} RMSE"] = mse[p, s].sqrt()
            logger.info(
                "{} {}\tMAE = {:0.3g}, RMSE = {:0.3g}".format(
                    place,
                    strain,
                    stats[f"{place} {strain} MAE"],
                    stats[f"{place} {strain} RMSE"],
                )
            )

    return {k: float(v) for k, v in stats.items()}

@torch.no_grad()
def log_holdout_stats(fits: dict) -> dict:
    """
    Logs statistics comparing multiple results from ``fit_svi``.
    """
    assert len(fits) > 1
    fits = list(fits.items())
    stats = {}
    for i, (name1, fit1) in enumerate(fits[:-1]):
        for name2, fit2 in fits[i + 1 :]:
            # Compute mutation similarity.
            mutations = sorted(set(fit1["mutations"]) & set(fit2["mutations"]))
            medians = []
            for fit in (fit1, fit2):
                mutation_id = {m: i for i, m in enumerate(fit["mutations"])}
                idx = torch.tensor([mutation_id[m] for m in mutations])
                medians.append(fit["median"]["coef"][idx] * 0.01)
            error = medians[0] - medians[1]
            mutation_std = torch.cat(medians).std().item()
            mutation_rmse = error.square().mean().sqrt().item()
            mutation_mae = error.abs().mean().item()
            mutation_correlation = pearson_correlation(medians[0], medians[1]).item()

            # Compute lineage similarity.
            means = []
            for fit in (fit1, fit2):
                rate = fit["mean"]["rate"]
                if rate.dim() == 2:
                    rate = rate.mean(0)
                means.append(rate)
            error = means[0] - means[1]
            lineage_std = torch.cat(means).std().item()
            lineage_rmse = error.square().mean().sqrt().item()
            lineage_mae = error.abs().mean().item()
            lineage_correlation = pearson_correlation(means[0], means[1]).item()

            # Print stats.
            logger.info(
                f"{name1} vs {name2} mutations: "
                f"ρ = {mutation_correlation:0.3g}, "
                f"RMSE = {mutation_rmse:0.3g}, "
                f"MAE = {mutation_mae:0.3g}"
            )
            logger.info(
                f"{name1} vs {name2} lineages: "
                f"ρ = {lineage_correlation:0.3g}, "
                f"RMSE = {lineage_rmse:0.3g}, "
                f"MAE = {lineage_mae:0.3g}"
            )

            # Save stats.
            stats["mutation_corr"] = mutation_correlation
            stats["mutation_rmse"] = mutation_rmse
            stats["mutation_mae"] = mutation_mae
            stats["mutation_stddev"] = mutation_std
            stats["lineage_corr"] = lineage_correlation
            stats["lineage_rmse"] = lineage_rmse
            stats["lineage_mae"] = lineage_mae
            stats["lineage_stdev"] = lineage_std

    return {k: float(v) for k, v in stats.items()}


# -*- coding: utf-8 -*-

import numpy as np

from pyfr.solvers.baseadvecdiff import (BaseAdvectionDiffusionBCInters,
                                        BaseAdvectionDiffusionIntInters,
                                        BaseAdvectionDiffusionMPIInters)

class NavierStokesIntInters(BaseAdvectionDiffusionIntInters):
    def __init__(self, *args, **kwargs):
        super(NavierStokesIntInters, self).__init__(*args, **kwargs)

        # Pointwise template arguments
        rsolver = self._cfg.get('solver-interfaces', 'riemann-solver')
        self._tplargs = dict(ndims=self.ndims, nvars=self.nvars,
                             rsolver=rsolver, c=self._tpl_c)

        self._be.pointwise.register('pyfr.solvers.navstokes.kernels.intconu')
        self._be.pointwise.register('pyfr.solvers.navstokes.kernels.intcflux')

    def get_con_u_kern(self):
        return self._be.kernel('intconu', self._tplargs,
                               dims=[self.ninterfpts],
                               ulin=self._scal0_lhs, urin=self._scal0_rhs,
                               ulout=self._scal1_lhs, urout=self._scal1_rhs)

    def get_comm_flux_kern(self):
        return self._be.kernel('intcflux', self._tplargs,
                               dims=[self.ninterfpts],
                               ul=self._scal0_lhs, ur=self._scal0_rhs,
                               gradul=self._vect0_lhs, gradur=self._vect0_rhs,
                               magnl=self._mag_pnorm_lhs,
                               magnr=self._mag_pnorm_rhs,
                               nl=self._norm_pnorm_lhs)

class NavierStokesMPIInters(BaseAdvectionDiffusionMPIInters):
    def __init__(self, *args, **kwargs):
        super(NavierStokesMPIInters, self).__init__(*args, **kwargs)

        # Pointwise template arguments
        rsolver = self._cfg.get('solver-interfaces', 'riemann-solver')
        self._tplargs = dict(ndims=self.ndims, nvars=self.nvars,
                             rsolver=rsolver, c=self._tpl_c)

        self._be.pointwise.register('pyfr.solvers.navstokes.kernels.mpiconu')
        self._be.pointwise.register('pyfr.solvers.navstokes.kernels.mpicflux')

    def get_con_u_kern(self):
        return self._be.kernel('mpiconu', self._tplargs,
                               dims=[self.ninterfpts],
                               ulin=self._scal0_lhs, urin=self._scal0_rhs,
                               ulout=self._scal1_lhs)

    def get_comm_flux_kern(self):
        return self._be.kernel('mpicflux', self._tplargs,
                               dims=[self.ninterfpts],
                               ul=self._scal0_lhs, ur=self._scal0_rhs,
                               gradul=self._vect0_lhs, gradur=self._vect0_rhs,
                               magnl=self._mag_pnorm_lhs,
                               nl=self._norm_pnorm_lhs)

class NavierStokesBaseBCInters(BaseAdvectionDiffusionBCInters):
    def __init__(self, *args, **kwargs):
        super(NavierStokesBaseBCInters, self).__init__(*args, **kwargs)

        # Pointwise template arguments
        rsolver = self._cfg.get('solver-interfaces', 'riemann-solver')
        self._tplargs = dict(ndims=self.ndims, nvars=self.nvars,
                             rsolver=rsolver, c=self._tpl_c,
                             bctype=self.type)

        self._be.pointwise.register('pyfr.solvers.navstokes.kernels.bcconu')
        self._be.pointwise.register('pyfr.solvers.navstokes.kernels.bccflux')

    def get_con_u_kern(self):
        return self._be.kernel('bcconu', self._tplargs, dims=[self.ninterfpts],
                               ulin=self._scal0_lhs, ulout=self._scal1_lhs)

    def get_comm_flux_kern(self):
        return self._be.kernel('bccflux', self._tplargs,
                               dims=[self.ninterfpts],
                               ul=self._scal0_lhs, gradul=self._vect0_lhs,
                               magnl=self._mag_pnorm_lhs,
                               nl=self._norm_pnorm_lhs)

class NavierStokesNoSlpIsotWallBCInters(NavierStokesBaseBCInters):
    type = 'no-slp-isot-wall'

    def __init__(self, *args, **kwargs):
        super(NavierStokesNoSlpIsotWallBCInters, self).__init__(*args,
                                                                **kwargs)

        self._tpl_c['cpTw'], = self._eval_opts(['cpTw'])
        self._tpl_c['v'] = self._eval_opts('uvw'[:self.ndims], default='0')

class NavierStokesNoSlpAdiaWallBCInters(NavierStokesBaseBCInters):
    type = 'no-slp-adia-wall'

class NavierStokesSupInflowBCInters(NavierStokesBaseBCInters):
    type = 'sup-in-fa'

    def __init__(self, *args, **kwargs):
        super(NavierStokesSupInflowBCInters, self).__init__(*args, **kwargs)

        self._tpl_c['rho'], self._tpl_c['p'] = self._eval_opts(['rho', 'p'])
        self._tpl_c['v'] = self._eval_opts('uvw'[:self.ndims])

class NavierStokesSupOutflowBCInters(NavierStokesBaseBCInters):
    type = 'sup-out-fn'

class NavierStokesSubInflowFrvBCInters(NavierStokesBaseBCInters):
    type = 'sub-in-frv'

    def __init__(self, *args, **kwargs):
        super(NavierStokesSubInflowFrvBCInters, self).__init__(*args, **kwargs)

        self._tpl_c['rho'], = self._eval_opts(['rho'])
        self._tpl_c['v'] = self._eval_opts('uvw'[:self.ndims])

class NavierStokesSubInflowFtpttangBCInters(NavierStokesBaseBCInters):
    type = 'sub-in-ftpttang'

    def __init__(self, *args, **kwargs):
        super(NavierStokesSubInflowFtpttangBCInters, self).__init__(*args,
                                                                    **kwargs)

        gamma = self._cfg.getfloat('constants', 'gamma')

        # Pass boundary constants to the backend
        self._tpl_c['cpTt'], = self._eval_opts(['cpTt'])
        self._tpl_c['pt'], = self._eval_opts(['pt'])
        self._tpl_c['Rdcp'] = (gamma - 1.0)/gamma

        # Calculate u, v velocity components from the inflow angle
        theta = self._eval_opts(['theta'])[0]*np.pi/180.0
        velcomps = np.array([np.cos(theta), np.sin(theta), 1.0])

        # Adjust u, v and calculate w velocity components for 3-D
        if self.ndims == 3:
            phi = self._eval_opts(['phi'])[0]*np.pi/180.0
            velcomps[:2] *= np.sin(phi)
            velcomps[2] *= np.cos(phi)

        self._tpl_c['vc'] = velcomps[:self.ndims]

class NavierStokesSubOutflowBCInters(NavierStokesBaseBCInters):
    type = 'sub-out-fp'

    def __init__(self, *args, **kwargs):
        super(NavierStokesSubOutflowBCInters, self).__init__(*args, **kwargs)

        self._tpl_c['p'], = self._eval_opts(['p'])