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

# Copyright (c) 2012, Sergio Callegari
# All rights reserved.

# This file is part of PyDSM.

# PyDSM is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# PyDSM is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with PyDSM.  If not, see <http://www.gnu.org/licenses/>.

"""
Computation of the impulse response of a DT filter (:mod:`pydsm.ir`)
====================================================================

Compute (approximating by truncation) the impulse response
of a discrete time filter, trying to (roughly) guess the appropriate
truncation.

.. currentmodule:: pydsm.ir

Functions
---------

.. autosummary::
   :toctree: generated/

    guess_ir_length  -- Guess appropriate truncation length
    impulse_response  -- Compute impulse response of DT filter
"""

from __future__ import division, print_function

import numpy as np
import scipy as sp
__import__("scipy.signal")

__all__ = ["impulse_response", "guess_ir_length"]

def impulse_response(h, m=None, db=80):
    """
    Computes the filter impulse response

    Parameters
    ----------
    h : tuple_like
        the filter definition either in zpk or in nd form.

    Returns
    -------
    ir : ndarray
        the truncated impulse response

    Other Parameters
    ----------------
    m : int, optional
        the number of samples after which the impulse response should be
        truncated. Defaults to None, which means *try to guess*
    db : real, optional
        a hint about how to guess the length where the impuls response
        should be truncated (defaults to 80)

    Notes
    -----
    The guess about the lenght where the impulse response can be truncated
    is extremely rough. See :func:`guess_ir_length` in this module for
    further info.
    """
    if len(h) == 3:
        (b, a) = sp.signal.zpk2tf(*h)
        b = b.real
        a = a.real
    else:
        (b, a) = h
    if m is None:
        m = guess_ir_length(h, db)
    ins = np.zeros(m)
    ins[0] = 1
    return sp.signal.lfilter(b, a, ins)

def guess_ir_length(h, db=80):
    """
    Tries to estimate an appropriate length for the filter response

    Parameters
    ----------
    h : tuple_like
        the filter definition either in zpk or in nd form.
    db : real, optional
        a hint about how to guess the length where the impulse response
        should be truncated. This is defined on a log scale. The larger
        the longer the resulting length. Defaults to 80.

    Returns
    -------
    m : int
        a guess about the appropriate number of samples to represent the
        filter impulse response with the required accuracy

    Notes
    -----
    The guess is based on the slowlest pole of the filter, considering when
    its response is attenuated to -db. This can be by far too optimistic in
    some cases and particularly when there are overlapping or similar poles.

    Do not try to use this function for filters with poles in 1.
    """
    # Put h in zpk form if it is in tf form
    if len(h) == 2:
        h = sp.signal.tf2zpk(*h)
    pp = h[1]
    t_z = len(h[0])+1
    if len(pp) == 0:
        t_p = 0
    else:
        # Try to estimate length of decay of the filter h.
        # The estimation is extremely rough, based on the decay
        # rate of the pole with maximum magnitude.
        # Thus, it breaks easily when there are poles very close
        # one to the other or overlapping.
        # Furthermore, this code should not be called if there is
        # a pole in z=1.
        os = np.seterr(divide='ignore')
        sr = np.log(np.abs(pp))
        np.seterr(**os)
        # Take slowlest pole
        wmin = np.min(np.abs(sr))
        # 1/omega min is time constant in sample periods.
        # Let's multiply the time constant in order to have
        # the transient attenuated by db decibels
        t_p = int(np.ceil(db/20*np.log(10)/wmin))
    return t_p+t_z


#-*-coding:utf-8-*-
'''
DpCas-Light
||||      ||||       ||||         ||        ||||
||  ||    ||  ||   ||    ||      ||||     ||    ||
||    ||  ||   || ||      ||    ||  ||     ||
||    ||  ||  ||  ||           ||====||      ||||
||    ||  ||||    ||      ||  ||======||         ||
||  ||    ||       ||    ||  ||        ||  ||    ||
||||      ||         ||||   ||          ||   ||||

/------------------ Who You Want 2 See ------------------/
'''
# date:2021-04-17
# Author: Eric.Lee
# function: pipline

import os
import numpy as np
import cv2
import torch
from PIL import Image

def compute_iou(rec1, rec2):
    """
    computing IoU
    :param rec1: (y0, x0, y1, x1), which reflects
            (top, left, bottom, right)
    :param rec2: (y0, x0, y1, x1)
    :return: scala value of IoU
    """
    # computing area of each rectangles
    S_rec1 = (rec1[2] - rec1[0]) * (rec1[3] - rec1[1])
    S_rec2 = (rec2[2] - rec2[0]) * (rec2[3] - rec2[1])

    # computing the sum_area
    sum_area = S_rec1 + S_rec2

    # find the each edge of intersect rectangle
    left_line = max(rec1[1], rec2[1])
    right_line = min(rec1[3], rec2[3])
    top_line = max(rec1[0], rec2[0])
    bottom_line = min(rec1[2], rec2[2])

    # judge if there is an intersect
    if left_line >= right_line or top_line >= bottom_line:
        return 0
    else:
        intersect = (right_line - left_line) * (bottom_line - top_line)
        #return (intersect / (sum_area - intersect))*1.0
        return (intersect / (S_rec1 + 1e-6))*1.0

def draw_landmarks(img,output,face_w,face_h,x0,y0,vis = False):
    img_width = img.shape[1]
    img_height = img.shape[0]
    dict_landmarks = {}
    eyes_center = []
    x_list = []
    y_list = []
    for i in range(int(output.shape[0]/2)):
        x = output[i*2+0]*float(face_w) + x0
        y = output[i*2+1]*float(face_h) + y0

        x_list.append(x)
        y_list.append(y)

        if 41>= i >=33:
            if 'left_eyebrow' not in dict_landmarks.keys():
                dict_landmarks['left_eyebrow'] = []
            dict_landmarks['left_eyebrow'].append([int(x),int(y),(0,255,0)])
            if vis:
                cv2.circle(img, (int(x),int(y)), 2, (0,255,0),-1)
        elif 50>= i >=42:
            if 'right_eyebrow' not in dict_landmarks.keys():
                dict_landmarks['right_eyebrow'] = []
            dict_landmarks['right_eyebrow'].append([int(x),int(y),(0,255,0)])
            if vis:
                cv2.circle(img, (int(x),int(y)), 2, (0,255,0),-1)
        elif 67>= i >=60:
            if 'left_eye' not in dict_landmarks.keys():
                dict_landmarks['left_eye'] = []
            dict_landmarks['left_eye'].append([int(x),int(y),(255,55,255)])
            if vis:
                cv2.circle(img, (int(x),int(y)), 2, (255,0,255),-1)
        elif 75>= i >=68:
            if 'right_eye' not in dict_landmarks.keys():
                dict_landmarks['right_eye'] = []
            dict_landmarks['right_eye'].append([int(x),int(y),(255,55,255)])
            if vis:
                cv2.circle(img, (int(x),int(y)), 2, (255,0,255),-1)
        elif 97>= i >=96:
            eyes_center.append((x,y))
            if vis:
                cv2.circle(img, (int(x),int(y)), 2, (0,0,255),-1)
        elif 54>= i >=51:
            if 'bridge_nose' not in dict_landmarks.keys():
                dict_landmarks['bridge_nose'] = []
            dict_landmarks['bridge_nose'].append([int(x),int(y),(0,170,255)])
            if vis:
                cv2.circle(img, (int(x),int(y)), 2, (0,170,255),-1)
        elif 32>= i >=0:
            if 'basin' not in dict_landmarks.keys():
                dict_landmarks['basin'] = []
            dict_landmarks['basin'].append([int(x),int(y),(255,30,30)])
            if vis:
                cv2.circle(img, (int(x),int(y)), 2, (255,30,30),-1)
        elif 59>= i >=55:
            if 'wing_nose' not in dict_landmarks.keys():
                dict_landmarks['wing_nose'] = []
            dict_landmarks['wing_nose'].append([int(x),int(y),(0,255,255)])
            if vis:
                cv2.circle(img, (int(x),int(y)), 2, (0,255,255),-1)
        elif 87>= i >=76:
            if 'out_lip' not in dict_landmarks.keys():
                dict_landmarks['out_lip'] = []
            dict_landmarks['out_lip'].append([int(x),int(y),(255,255,0)])
            if vis:
                cv2.circle(img, (int(x),int(y)), 2, (255,255,0),-1)
        elif 95>= i >=88:
            if 'in_lip' not in dict_landmarks.keys():
                dict_landmarks['in_lip'] = []
            dict_landmarks['in_lip'].append([int(x),int(y),(50,220,255)])
            if vis:
                cv2.circle(img, (int(x),int(y)), 2, (50,220,255),-1)
        else:
            if vis:
                cv2.circle(img, (int(x),int(y)), 2, (255,0,255),-1)
    face_area = (max(x_list) - min(x_list))*(max(y_list) - min(y_list))
    return dict_landmarks,eyes_center,face_area

def draw_contour(image,dict,vis = False):
    x0 = 0# 偏置
    y0 = 0

    for key in dict.keys():
        # print(key)
        _,_,color = dict[key][0]

        if 'left_eye' == key:
            eye_x = np.mean([dict[key][i][0]+x0 for i in range(len(dict[key]))])
            eye_y = np.mean([dict[key][i][1]+y0 for i in range(len(dict[key]))])
            if vis:
                cv2.circle(image, (int(eye_x),int(eye_y)), 3, (255,255,55),-1)
        if 'right_eye' == key:
            eye_x = np.mean([dict[key][i][0]+x0 for i in range(len(dict[key]))])
            eye_y = np.mean([dict[key][i][1]+y0 for i in range(len(dict[key]))])
            if vis:
                cv2.circle(image, (int(eye_x),int(eye_y)), 3, (255,215,25),-1)

        if 'basin' == key or 'wing_nose' == key:
            pts = np.array([[dict[key][i][0]+x0,dict[key][i][1]+y0] for i in range(len(dict[key]))],np.int32)
            if vis:
                cv2.polylines(image,[pts],False,color,thickness = 2)

        else:
            points_array = np.zeros((1,len(dict[key]),2),dtype = np.int32)
            for i in range(len(dict[key])):
                x,y,_ = dict[key][i]
                points_array[0,i,0] = x+x0
                points_array[0,i,1] = y+y0

            # cv2.fillPoly(image, points_array, color)
            if vis:
                cv2.drawContours(image,points_array,-1,color,thickness=2)

def plot_box(x, img, color=None, label=None, line_thickness=None):
    # Plots one bounding box on image img
    tl = line_thickness or round(0.002 * max(img.shape[0:2])) + 1  # line thickness
    color = color or [random.randint(0, 255) for _ in range(3)]
    c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
    cv2.rectangle(img, c1, c2, color, thickness=tl)
    if label:
        tf = max(tl - 1, 2)  # font thickness
        t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
        c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
        cv2.rectangle(img, c1, c2, color, -1)  # filled
        cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [185, 195,190], thickness=tf, lineType=cv2.LINE_AA)
#-------------------------------------------------------------------------------

def face_alignment(imgn,eye_left_n,eye_right_n,\
desiredLeftEye=(0.34, 0.42),desiredFaceWidth=256, desiredFaceHeight=None):

    if desiredFaceHeight is None:
        desiredFaceHeight = desiredFaceWidth

    leftEyeCenter = eye_left_n
    rightEyeCenter = eye_right_n
    # compute the angle between the eye centroids
    dY = rightEyeCenter[1] - leftEyeCenter[1]
    dX = rightEyeCenter[0] - leftEyeCenter[0]
    angle = np.degrees(np.arctan2(dY, dX))

    # compute the desired right eye x-coordinate based on the
    # desired x-coordinate of the left eye
    desiredRightEyeX = 1.0 - desiredLeftEye[0]
	# determine the scale of the new resulting image by taking
	# the ratio of the distance between eyes in the *current*
	# image to the ratio of distance between eyes in the
	# *desired* image
    dist = np.sqrt((dX ** 2) + (dY ** 2))
    desiredDist = (desiredRightEyeX - desiredLeftEye[0])
    desiredDist *= desiredFaceWidth
    scale = desiredDist / dist
    # compute center (x, y)-coordinates (i.e., the median point)
    # between the two eyes in the input image
    eyesCenter = ((leftEyeCenter[0] + rightEyeCenter[0]) / 2,(leftEyeCenter[1] + rightEyeCenter[1]) / 2)
    # grab the rotation matrix for rotating and scaling the face
    M = cv2.getRotationMatrix2D(eyesCenter, angle, scale)
    # update the translation component of the matrix
    tX = desiredFaceWidth * 0.5
    tY = desiredFaceHeight * desiredLeftEye[1]
    M[0, 2] += (tX - eyesCenter[0])
    M[1, 2] += (tY - eyesCenter[1])

    M_reg = np.zeros((3,3),dtype = np.float32)
    M_reg[0,:] = M[0,:]
    M_reg[1,:] = M[1,:]
    M_reg[2,:] = (0,0,1.)
    # print(M_reg)
    M_I = np.linalg.inv(M_reg)#矩阵求逆，从而获得，目标图到原图的关系
    # print(M_I)
    # apply the affine transformation
    (w, h) = (desiredFaceWidth, desiredFaceHeight)
    # cv_resize_model = [cv2.INTER_LINEAR,cv2.INTER_CUBIC,cv2.INTER_NEAREST,cv2.INTER_AREA]

    output = cv2.warpAffine(imgn, M, (w, h),flags=cv2.INTER_LINEAR,borderMode=cv2.BORDER_CONSTANT)#

    #---------------------------------------------------------------------------------------

    # ptx1 = int(eye_left_gt_n[0]*M[0][0] + eye_left_gt_n[1]*M[0][1] + M[0][2])
    # pty1 = int(eye_left_gt_n[0]*M[1][0] + eye_left_gt_n[1]*M[1][1] + M[1][2])
    #
    # ptx2 = int(eye_right_gt_n[0]*M[0][0] + eye_right_gt_n[1]*M[0][1] + M[0][2])
    # pty2 = int(eye_right_gt_n[0]*M[1][0] + eye_right_gt_n[1]*M[1][1] + M[1][2])

    return output

def refine_face_bbox(bbox,img_shape):
    height,width,_ = img_shape

    x1,y1,x2,y2 = bbox

    expand_w = (x2-x1)
    expand_h = (y2-y1)

    x1 -= expand_w*0.12
    y1 -= expand_h*0.12
    x2 += expand_w*0.12
    y2 += expand_h*0.08

    x1,y1,x2,y2 = int(x1),int(y1),int(x2),int(y2)

    x1 = np.clip(x1,0,width-1)
    y1 = np.clip(y1,0,height-1)
    x2 = np.clip(x2,0,width-1)
    y2 = np.clip(y2,0,height-1)

    return (x1,y1,x2,y2)

def get_faces_batch_attribute(face_multitask_model,face_euler_model,dets,img_raw,use_cuda,face_size = 256,vis = False):

    face_map = np.zeros([112*3,112*3,3]).astype(np.uint8)
    face_map[:,:,0].fill(205)
    face_map[:,:,1].fill(205)
    face_map[:,:,2].fill(205)
    if len(dets) == 0:
        return [],[],[],face_map
    img_align = img_raw.copy()
    # 绘制图像
    image_batch = None
    r_bboxes = []
    imgs_crop = []
    for b in dets:
        b = list(map(int, b))

        r_bbox = refine_face_bbox((b[0],b[1],b[2],b[3]),img_raw.shape)
        r_bboxes.append(r_bbox)
        img_crop = img_raw[r_bbox[1]:r_bbox[3],r_bbox[0]:r_bbox[2]]
        imgs_crop.append(img_crop)
        img_ = cv2.resize(img_crop, (face_size,face_size), interpolation = cv2.INTER_LINEAR) # INTER_LINEAR INTER_CUBIC

        img_ = img_.astype(np.float32)
        img_ = (img_-128.)/256.

        img_ = img_.transpose(2, 0, 1)
        img_ = np.expand_dims(img_,0)

        if image_batch is None:
            image_batch = img_
        else:
            image_batch = np.concatenate((image_batch,img_),axis=0)

    # # 填充最大 关键点 批次数据
    # if len(dets) < ops.max_batch_size:
    #     im_mask = np.zeros([1,3,ops.landmarks_img_size[0],ops.landmarks_img_size[1]], dtype = np.float32)
    #     for i in range(ops.max_batch_size-len(dets)):
    #         if image_batch is None:
    #             image_batch = im_mask
    #         else:
    #             image_batch = np.concatenate((image_batch,im_mask),axis=0)
    #
    # print("image_batch shape:",image_batch.shape)
    # image_batch = torch.from_numpy(image_batch).float()
    # #
    # if use_cuda:
    #     image_batch = image_batch.cuda()  # (bs, 3, h, w)

    landmarks_pre,gender_pre,age_pre = face_multitask_model.predict(image_batch)
    euler_angles = face_euler_model.predict(image_batch)
    # print(" -------->>> euler_angles : ",euler_angles)
    # print("landmarks_pre,gender_pre,age_pre :",landmarks_pre.shape,gender_pre.shape,age_pre.shape)

    faces_identify = [] # 符合要求，需要识别的人的图像
    faces_identify_bboxes = []# 符合要求，需要识别的人脸边界框
    faceid_idx = 0 # 符合要求，需要识别的人的索引计数器
    for i in range(len(dets)):
        x0,y0 = r_bboxes[i][0],r_bboxes[i][1]
        face_w = r_bboxes[i][2]-r_bboxes[i][0]
        face_h = r_bboxes[i][3]-r_bboxes[i][1]
        dict_landmarks,eyes_center,face_area = draw_landmarks(img_raw,landmarks_pre[i],face_w,face_h,x0,y0,vis = False)

        gray_ = cv2.cvtColor(img_align[r_bboxes[i][1]:r_bboxes[i][3],r_bboxes[i][0]:r_bboxes[i][2],:], cv2.COLOR_BGR2GRAY)

        blur_ = cv2.Laplacian(gray_, cv2.CV_64F).var()

        gender_max_index = np.argmax(gender_pre[i])#概率最大类别索引
        score_gender = gender_pre[i][gender_max_index]# 最大概率

        yaw,pitch,roll = euler_angles[i]

        cv2.putText(img_raw, "yaw:{:.1f},pitch:{:.1f},roll:{:.1f}".format(yaw,pitch,roll),(int(r_bboxes[i][0]-20),int(r_bboxes[i][1]-30)),cv2.FONT_HERSHEY_DUPLEX, 0.65, (253,139,54), 5)
        cv2.putText(img_raw, "yaw:{:.1f},pitch:{:.1f},roll:{:.1f}".format(yaw,pitch,roll),(int(r_bboxes[i][0]-20),int(r_bboxes[i][1]-30)),cv2.FONT_HERSHEY_DUPLEX, 0.65, (20,185,255), 1)

        cv2.putText(img_raw, "{}".format(int(face_area)),(int(r_bboxes[i][0]-1),int(r_bboxes[i][3]-3)),cv2.FONT_HERSHEY_DUPLEX, 0.65, (253,39,54), 5) # face_area
        cv2.putText(img_raw, "{}".format(int(face_area)),(int(r_bboxes[i][0]-1),int(r_bboxes[i][3]-3)),cv2.FONT_HERSHEY_DUPLEX, 0.65, (20,185,255), 1)
        if gender_max_index == 1.:
            gender_str = "male"
        else:
            gender_str = "female"

        if abs(yaw)<45.:
            face_align_output = face_alignment(img_align,eyes_center[0],eyes_center[1],
                desiredLeftEye=(0.365, 0.38),desiredFaceWidth=112, desiredFaceHeight=None)
        else:
            face_align_output = face_alignment(img_align,eyes_center[0],eyes_center[1],
                desiredLeftEye=(0.38, 0.40),desiredFaceWidth=112, desiredFaceHeight=None)

        # plot_box(r_bboxes[i][0:4], img_raw,label="{}, age: {:.1f}, unblur:{}".format(gender_str,age_pre[i][0],int(blur_)), color=(255,90,90), line_thickness=2)
        plot_box(r_bboxes[i][0:4], img_raw,label="{}, age: {:.1f}".format(gender_str,age_pre[i][0]), color=(255,90,90), line_thickness=2)
        # print("face_area:",face_area)
        # if (blur_>35) and abs(yaw)<36. and abs(pitch)<30. and (face_area>(60*60)):
        if abs(yaw)<36. and abs(pitch)<36. and (face_area>(60*60)):
            if vis :
                draw_contour(img_raw,dict_landmarks,vis = True)

            faces_identify.append(Image.fromarray(face_align_output))
            faces_identify_bboxes.append(r_bboxes[i][0:4])

            if faceid_idx<9:
                y1_map,y2_map = int(faceid_idx/3)*112,(int(faceid_idx/3)+1)*112
                x1_map,x2_map = int(faceid_idx%3)*112,(int(faceid_idx%3)+1)*112
                face_map[y1_map:y2_map,x1_map:x2_map,:] = face_align_output
                cv2.rectangle(face_map, (int(x1_map),int(y1_map)), (int(x2_map),int(y2_map)), (55,255,255), 2)
                faceid_idx += 1
        else:
            cv2.putText(img_raw, "bad for face reco",(int(r_bboxes[i][0]-1),int(r_bboxes[i][3]+20)),cv2.FONT_HERSHEY_DUPLEX, 0.65, (20,15,255), 4)
            cv2.putText(img_raw, "bad for face reco",(int(r_bboxes[i][0]-1),int(r_bboxes[i][3]+20)),cv2.FONT_HERSHEY_DUPLEX, 0.65, (220,185,25), 1)

    return faces_identify,faces_identify_bboxes,r_bboxes,face_map


#!/usr/bin/env python
"""
substitution_model.py

Contains classes for defining Markov models of substitution.
These classes depend on an Alphabet class member for defining the set
of motifs that each represent a state in the Markov chain. Examples of
a 'dna' type alphabet motif is 'a', and of a 'codon' type motif is'atg'.

By default all models include the gap motif ('-' for a 'dna' alphabet or
'---' for a 'codon' alphabet). This differs from software such as PAML,
where gaps are treated as ambiguituous states (specifically, as 'n'). The gap
motif state can be excluded from the substitution model using the method
excludeGapMotif(). It is recommended that to ensure the alignment and the
substitution model are defined with the same alphabet that modifications
are done to the substitution model alphabet and this instance is then given
to the alignment.

The model's substitution rate parameters are represented as a dictionary
with the parameter names as keys, and predicate functions as the values.
These predicate functions compare a pair of motifs, returning True or False.
Many such functions are provided as methods of the class. For instance,
the istransition method is pertinent to dna based models. This method returns
True if an 'a'/'g' or 'c'/'t' pair is passed to it, False otherwise. In this
way the positioning of parameters in the instantaneous rate matrix (commonly
called Q) is determined.

>>> model = Nucleotide(equal_motif_probs=True)
>>> model.setparameterrules({'alpha': model.istransition})
>>> parameter_controller = model.make_likelihood_function(tree)
"""

import inspect
import json
import warnings

from collections.abc import Callable

import numpy

from numpy.linalg import svd

from cogent3.core import moltype
from cogent3.evolve import motif_prob_model, parameter_controller, predicate
from cogent3.evolve.discrete_markov import PsubMatrixDefn
from cogent3.evolve.likelihood_tree import make_likelihood_tree_leaf
from cogent3.evolve.substitution_calculation import (
    AlignmentAdaptDefn,
    CalcDefn,
    CallDefn,
    ConstDefn,
    ExpDefn,
    GammaDefn,
    LengthDefn,
    MonotonicDefn,
    NonParamDefn,
    PartitionDefn,
    ProductDefn,
    RateDefn,
    SelectForDimension,
)
from cogent3.evolve.substitution_calculation import (
    SubstitutionParameterDefn as ParamDefn,
)
from cogent3.evolve.substitution_calculation import WeightedPartitionDefn
from cogent3.maths.optimisers import ParameterOutOfBoundsError
from cogent3.util.misc import extend_docstring_from, get_object_provenance

__author__ = "Peter Maxwell, Gavin Huttley and Andrew Butterfield"
__copyright__ = "Copyright 2007-2020, The Cogent Project"
__contributors__ = [
    "Gavin Huttley",
    "Andrew Butterfield",
    "Peter Maxwell",
    "Matthew Wakefield",
    "Brett Easton",
    "Rob Knight",
    "Von Bing Yap",
]
__license__ = "BSD-3"
__version__ = "2020.7.2a"
__maintainer__ = "Gavin Huttley"
__email__ = "gavin.huttley@anu.edu.au"
__status__ = "Production"

def predicate2matrix(alphabet, pred, mask=None):
    """From a test like istransition() produce an MxM boolean matrix"""
    M = len(alphabet)
    result = numpy.zeros([M, M], int)
    for i in range(M):
        for j in range(M):
            if mask is None or mask[i, j]:
                result[i, j] = pred(alphabet[i], alphabet[j])
    return result

def redundancy_in_predicate_masks(preds):
    # Calculate the nullity of the predicates.  If non-zero
    # there is some redundancy and the model will be overparameterised.
    if len(preds) <= 1:
        return 0
    eqns = 1.0 * numpy.array([list(mask.flat) for mask in list(preds.values())])
    svs = svd(eqns)[1]
    # count non-duplicate non-zeros singular values
    matrix_rank = len([sv for sv in svs if abs(sv) > 1e-8])
    return len(preds) - matrix_rank

def _maxWidthIfTruncated(pars, delim, each):
    # 'pars' is an array of lists of strings, how long would the longest
    # list representation be if the strings were truncated at 'each'
    # characters and joined together with 'delim'.
    return max(
        [
            sum([min(len(par), each) for par in par_list])
            + len(delim) * (len(par_list) - 1)
            for par_list in pars.flat
        ]
    )

def _isSymmetrical(matrix):
    return numpy.alltrue(numpy.alltrue(matrix == numpy.transpose(matrix)))

class _SubstitutionModel(object):
    # Subclasses must provide
    #  .make_param_controller_defns()

    def __init__(
        self,
        alphabet,
        motif_probs=None,
        optimise_motif_probs=False,
        equal_motif_probs=False,
        motif_probs_from_data=None,
        motif_probs_alignment=None,
        mprob_model=None,
        model_gaps=False,
        recode_gaps=False,
        motif_length=1,
        name="",
        motifs=None,
    ):
        # subclasses can extend this incomplete docstring
        """

        alphabet:
         - alphabet - An Alphabet object
         - motif_length: Use a tuple alphabet based on 'alphabet'.
         - motifs: Use a subalphabet that only contains those motifs.
         - model_gaps: Whether the gap motif should be included as a state.
         - recode_gaps: Whether gaps in an alignment should be treated as an
           ambiguous state instead.

        Motif Probability:
         - motif_probs: Dictionary of probabilities.
         - equal_motif_probs: Flag to set alignment motif probs equal.
         - motif_probs_alignment: An alignment from which motif probs are set.

         If none of these options are set then motif probs will be derived
         from the data: ie the particular alignment provided later.

         - optimise_motif_probs: Treat like other free parameters.  Any values
           set by the other motif_prob options will be used as initial values.

         - mprob_model: 'tuple', 'conditional', 'monomer' or 'monomers' to specify how
           tuple-alphabet (including codon) motif probs are used.

        """
        d = locals()
        exclude = ("self", "__class__")
        self._serialisable = {k: v for k, v in d.items() if k not in exclude}
        # MISC
        assert len(alphabet) < 65, (
            "Alphabet too big. Try explicitly " "setting alphabet to PROTEIN or DNA"
        )

        self.name = name
        self._optimise_motif_probs = optimise_motif_probs

        # ALPHABET
        if recode_gaps:
            if model_gaps:
                warnings.warn("Converting gaps to wildcards AND modeling gaps")
            else:
                model_gaps = False

        self.recode_gaps = recode_gaps

        self.moltype = alphabet.moltype
        if model_gaps:
            alphabet = alphabet.with_gap_motif()

        if motif_length > 1:
            alphabet = alphabet.get_word_alphabet(motif_length)

        if motifs is not None:
            alphabet = alphabet.get_subset(motifs)
        self.alphabet = alphabet
        self.gapmotif = alphabet.get_gap_motif()
        self._word_length = alphabet.get_motif_len()

        # MOTIF PROB ALPHABET MAPPING
        if mprob_model is None:
            mprob_model = "tuple" if self._word_length == 1 else "conditional"
        elif mprob_model == "word":
            mprob_model = "tuple"

        if model_gaps and mprob_model != "tuple":
            raise ValueError("mprob_model must be 'tuple' to model gaps")

        isinst = self._is_instantaneous
        self._instantaneous_mask = predicate2matrix(self.alphabet, isinst)
        self._instantaneous_mask_f = self._instantaneous_mask * 1.0
        self._mprob_model = mprob_model
        self.mprob_model = motif_prob_model.make_model(
            mprob_model, alphabet, self._instantaneous_mask_f
        )

        # MOTIF PROBS
        if equal_motif_probs:
            assert not (
                motif_probs or motif_probs_alignment
            ), "Motif probs equal or provided but not both"
            motif_probs = self.mprob_model.make_equal_motif_probs()
        elif motif_probs_alignment is not None:
            assert (
                not motif_probs
            ), "Motif probs from alignment or provided but not both"
            motif_probs = self.count_motifs(motif_probs_alignment)
            motif_probs = motif_probs.astype(float) / sum(motif_probs)
            assert len(alphabet) == len(motif_probs)
            motif_probs = dict(list(zip(alphabet, motif_probs)))
        if motif_probs:
            self.adapt_motif_probs(motif_probs)  # to check
            self.motif_probs = motif_probs
            if motif_probs_from_data is None:
                motif_probs_from_data = False
        else:
            self.motif_probs = None
            if motif_probs_from_data is None:
                motif_probs_from_data = True
        self.motif_probs_from_align = motif_probs_from_data

    def __getnewargs_ex__(self, *args, **kw):
        data = self.to_rich_dict(for_pickle=True)
        return (), data

    def to_rich_dict(self, for_pickle=False):
        data = self._serialisable.copy()
        if not for_pickle:
            for key, value in data.items():
                type_ = get_object_provenance(value)
                if type_.startswith("cogent3"):
                    try:
                        value = value.to_rich_dict(for_pickle=False)
                    except AttributeError:
                        pass
                    finally:
                        data[key] = value
            if "predicates" in data and data["predicates"]:
                data["predicates"] = [str(p) for p in data["predicates"]]
            data["type"] = get_object_provenance(self)
            data["version"] = __version__
        return data

    def to_json(self):
        """returns result of json formatted string"""
        data = self.to_rich_dict(for_pickle=False)
        return json.dumps(data)

    def get_param_list(self):
        return []

    def __str__(self):
        s = ["\n%s (" % self.__class__.__name__]
        s.append(
            "name = '%s'; type = '%s';"
            % (getattr(self, "name", None), getattr(self, "type", None))
        )
        if hasattr(self, "predicate_masks"):
            parlist = list(self.predicate_masks.keys())
            s.append("params = %s;" % parlist)
        motifs = self.get_motifs()
        s.append("number of motifs = %s;" % len(motifs))
        s.append("motifs = %s)\n" % motifs)
        return " ".join(s)

    def get_alphabet(self):
        return self.alphabet

    def get_mprob_alphabet(self):
        return self.mprob_model.get_input_alphabet()

    def get_motifs(self):
        return list(self.get_alphabet())

    @property
    def word_length(self):
        return self._word_length

    def get_motif_probs(self):
        """Return the dictionary of motif probabilities."""
        return self.motif_probs.copy()

    def set_param_controller_motif_probs(self, pc, mprobs, **kw):
        return self.mprob_model.set_param_controller_motif_probs(pc, mprobs, **kw)

    def make_likelihood_function(
        self,
        tree,
        motif_probs_from_align=None,
        optimise_motif_probs=None,
        aligned=True,
        expm=None,
        digits=None,
        space=None,
        **kw,
    ):

        if motif_probs_from_align is None:
            motif_probs_from_align = self.motif_probs_from_align

        if optimise_motif_probs is None:
            optimise_motif_probs = self._optimise_motif_probs

        kw["optimise_motif_probs"] = optimise_motif_probs
        kw["motif_probs_from_align"] = motif_probs_from_align

        if aligned:
            klass = parameter_controller.AlignmentLikelihoodFunction
        else:
            alphabet = self.get_alphabet()
            assert alphabet.get_gap_motif() not in alphabet
            klass = parameter_controller.SequenceLikelihoodFunction

        result = klass(self, tree, **kw)

        if self.motif_probs is not None:
            result.set_motif_probs(
                self.motif_probs, is_constant=not optimise_motif_probs, auto=True
            )

        if expm is None:
            expm = self._default_expm_setting
        if expm is not None:
            result.set_expm(expm)

        if digits or space:
            result.set_tables_format(digits=digits, space=space)

        return result

    def convert_alignment(self, alignment):
        # this is to support for everything but HMM
        result = {}
        for seq_name in alignment.names:
            sequence = alignment.get_gapped_seq(seq_name, self.recode_gaps)
            result[seq_name] = self.convert_sequence(sequence, seq_name)
        return result

    def convert_sequence(self, sequence, name):
        # make_likelihood_tree_leaf, sort of an indexed profile where duplicate
        # columns stored once, so likelihoods only calc'd once
        return make_likelihood_tree_leaf(sequence, self.get_alphabet(), name)

    def count_motifs(self, alignment, include_ambiguity=False):
        return self.mprob_model.count_motifs(
            alignment, include_ambiguity, self.recode_gaps
        )

    def make_alignment_defn(self, model):
        align = NonParamDefn("alignment", ("locus",))
        # The name of this matters, it's used in likelihood_function.py
        # to retrieve the correct (adapted) alignment.
        return AlignmentAdaptDefn(model, align)

    def adapt_motif_probs(self, motif_probs, auto=False):
        return self.mprob_model.adapt_motif_probs(motif_probs, auto=auto)

    def calc_monomer_probs(self, word_probs):
        # Not presently used, always go monomer->word instead
        return self.mprob_model.calc_monomer_probs(word_probs)

    def calc_word_probs(self, monomer_probs):
        return self.mprob_model.calc_word_probs(monomer_probs)

    def calc_word_weight_matrix(self, monomer_probs):
        return self.mprob_model.calc_word_weight_matrix(monomer_probs)

    def make_param_controller_defns(self, bin_names, endAtQd=False):
        (
            input_probs,
            word_probs,
            mprobs_matrix,
        ) = self.mprob_model.make_motif_word_prob_defns()

        if len(bin_names) > 1:
            bprobs = PartitionDefn(
                [1.0 / len(bin_names) for bin in bin_names],
                name="bprobs",
                dimensions=["locus"],
                dimension=("bin", bin_names),
            )
        else:
            bprobs = None

        defns = {
            "align": self.make_alignment_defn(ConstDefn(self, "model")),
            "bprobs": bprobs,
            "word_probs": word_probs,
        }

        rate_params = self.make_rate_params(bprobs)
        if endAtQd:
            defns["Qd"] = self.make_Qd_defn(word_probs, mprobs_matrix, rate_params)
        else:
            defns["psubs"] = self.make_psubs_defn(
                bprobs, word_probs, mprobs_matrix, rate_params
            )
        return defns

def non_zero_coords(matrix):
    dim = matrix.shape[0]
    coords = [(i, j) for i in range(dim) for j in range(dim) if matrix[i, j] != 0]
    return coords

class _ContinuousSubstitutionModel(_SubstitutionModel):
    # subclass must provide:
    #
    # - parameter_order: a list of parameter names corresponding to the
    #   arguments of:
    #
    # - calc_exchangeability_matrix(*params)
    #   convert len(self.parameter_order) params to a matrix

    """A substitution model for which the rate matrix (P) is derived from an
    instantaneous rate matrix (Q).  The nature of the parameters used to define
    Q is up to the subclasses.
    """

    # At some point this can be made variable, and probably
    # the default changed to False
    long_indels_are_instantaneous = True

    _exponentiator = None
    _default_expm_setting = "either"

    @extend_docstring_from(_SubstitutionModel.__init__)
    def __init__(
        self,
        alphabet,
        with_rate=False,
        ordered_param=None,
        distribution=None,
        partitioned_params=None,
        **kw,
    ):
        """
        - with_rate: Add a 'rate' parameter which varies by bin.
        - ordered_param: name of a single parameter which distinguishes any bins.
        - distribution: choices of 'free' or 'gamma' or an instance of some
          distribution. Could probably just deprecate free
        - partitioned_params: names of params to be partitioned across bins
        """

        _SubstitutionModel.__init__(self, alphabet, **kw)
        d = locals()
        exclude = ("self", "__class__")
        d = {k: v for k, v in d.items() if k not in exclude}
        self._serialisable.update(d)
        alphabet = self.get_alphabet()  # as may be altered by recode_gaps etc.

        # BINS
        if not ordered_param:
            if ordered_param is not None:
                warnings.warn("ordered_param should be a string or None")
                ordered_param = None
            if distribution:
                if with_rate:
                    ordered_param = "rate"
                else:
                    raise ValueError("distribution provided without ordered_param")
        elif not isinstance(ordered_param, str):
            warnings.warn("ordered_param should be a string or None")
            assert len(ordered_param) == 1, "More than one ordered_param"
            ordered_param = ordered_param[0]
            assert ordered_param, "False value hidden in list"
        self.ordered_param = ordered_param

        if distribution == "gamma":
            distribution = GammaDefn
        elif distribution in [None, "free"]:
            distribution = MonotonicDefn
        elif isinstance(distribution, str):
            raise ValueError('Unknown distribution "%s"' % distribution)
        self.distrib_class = distribution

        if not partitioned_params:
            partitioned_params = ()
        elif isinstance(partitioned_params, str):
            partitioned_params = (partitioned_params,)
        else:
            partitioned_params = tuple(partitioned_params)
        if self.ordered_param:
            if self.ordered_param not in partitioned_params:
                partitioned_params += (self.ordered_param,)
        self.partitioned_params = partitioned_params

        if "rate" in partitioned_params:
            with_rate = True
        self.with_rate = with_rate

        # CACHED SHORTCUTS
        self._exponentiator = None
        # self._ident = numpy.identity(len(self.alphabet), float)

    def check_params_exist(self):
        """Raise an error if the parameters specified to be partitioned or
        ordered don't actually exist."""
        for param in self.partitioned_params:
            if param not in self.parameter_order and param != "rate":
                desc = ["partitioned", "ordered"][param == self.ordered_param]
                raise ValueError('%s param "%s" unknown' % (desc, param))

    def _is_instantaneous(self, x, y):
        diffs = sum([X != Y for (X, Y) in zip(x, y)])
        return diffs == 1 or (
            diffs > 1
            and self.long_indels_are_instantaneous
            and self._is_any_indel(x, y)
        )

    def _is_any_indel(self, x, y):
        """An indel of any length"""
        # Things get complicated when a contigous indel of any length is OK:
        if x == y:
            return False
        gap_start = gap_end = gap_strand = None
        for (i, (X, Y)) in enumerate(zip(x, y)):
            G = self.gapmotif[i]
            if X != Y:
                if X != G and Y != G:
                    return False  # non-gap differences had their chance above
                elif gap_start is None:
                    gap_start = i
                    gap_strand = [X, Y].index(G)
                elif gap_end is not None or [X, Y].index(G) != gap_strand:
                    return False  # can't start a second gap
                else:
                    pass  # extend open gap
            elif gap_start is not None:
                gap_end = i
        return True

    def calcQ(self, word_probs, mprobs_matrix, *params):
        Q = self.calc_exchangeability_matrix(word_probs, *params)
        row_totals = Q.sum(axis=1)
        Q -= numpy.diag(row_totals)
        Q *= 1.0 / (word_probs * row_totals).sum()
        return Q

    def get_reference_cell(self):
        """returns the reference cell of a given model"""
        dim = len(self.alphabet)
        mats = numpy.zeros((dim, dim), dtype=int)
        for m in self.predicate_masks.values():
            mats += m
        ref_mask = self._instantaneous_mask - mats
        ref_cells = set(non_zero_coords(ref_mask))
        return ref_cells

    def get_param_matrix_coords(self, include_ref_cell=False):
        """returncoordinates for every predicate"""
        dim = len(self.alphabet)
        mats = numpy.zeros((dim, dim), dtype=int)
        param_coords = {}
        for key, m in self.predicate_masks.items():
            coords = [(i, j) for i in range(dim) for j in range(dim) if m[i, j] != 0]
            coords = set(coords)
            param_coords[key] = coords

        if include_ref_cell:
            param_coords["ref_cell"] = self.get_reference_cell()
        return param_coords

    def make_Qd_defn(self, word_probs, mprobs_matrix, rate_params):
        """Diagonalized Q, ie: rate matrix prepared for exponentiation"""
        Q = CalcDefn(self.calcQ, name="Q")(word_probs, mprobs_matrix, *rate_params)
        expm = NonParamDefn("expm")
        exp = ExpDefn(expm)
        Qd = CallDefn(exp, Q, name="Qd")
        return Qd

    def _make_bin_param_defn(self, edge_par_name, bin_par_name, bprob_defn):
        # if no ordered param defined, behaves as old, everything indexed by
        # and edge
        if edge_par_name not in self.partitioned_params:
            return ParamDefn(dimensions=["bin"], name=bin_par_name)

        if edge_par_name == self.ordered_param:
            whole = self.distrib_class(bprob_defn, bin_par_name)
        else:
            # this forces them to average to one, but no forced order
            # this means you can't force a param value to be shared across bins
            # so 1st above approach has to be used
            whole = WeightedPartitionDefn(bprob_defn, bin_par_name + "_partn")
        whole.bin_names = bprob_defn.bin_names
        return SelectForDimension(whole, "bin", name=bin_par_name)

    def make_rate_params(self, bprobs):
        params = []
        for param_name in self.parameter_order:
            if bprobs is None or param_name not in self.partitioned_params:
                defn = ParamDefn(param_name)
            else:
                e_defn = ParamDefn(param_name, dimensions=["edge", "locus"])
                # should be weighted by bprobs*rates not bprobs
                b_defn = self._make_bin_param_defn(
                    param_name, param_name + "_factor", bprobs
                )
                defn = ProductDefn(b_defn, e_defn, name=param_name + "_BE")
            params.append(defn)
        return params

    def make_fundamental_param_controller_defns(self, bin_names):
        """Everything one step short of the psubs, because cogent3.align code
        needs to handle Q*t itself."""
        defns = self.make_param_controller_defns(bin_names, endAtQd=True)
        assert "length" not in defns
        defns["length"] = LengthDefn()
        return defns

    def make_psubs_defn(self, bprobs, word_probs, mprobs_matrix, rate_params):
        distance = self.make_distance_defn(bprobs)
        P = self.make_continuous_psub_defn(
            word_probs, mprobs_matrix, distance, rate_params
        )
        return P

    def make_distance_defn(self, bprobs):
        length = LengthDefn()
        if self.with_rate and bprobs is not None:
            b_rate = self._make_bin_param_defn("rate", "rate", bprobs)
            distance = ProductDefn(length, b_rate, name="distance")
        else:
            distance = length
        return distance

    def make_continuous_psub_defn(
        self, word_probs, mprobs_matrix, distance, rate_params
    ):
        Qd = self.make_Qd_defn(word_probs, mprobs_matrix, rate_params)
        P = CallDefn(Qd, distance, name="psubs")
        return P

class StationaryQ:
    "Contains the Original Definition of calcQ"

    def calcQ(self, word_probs, mprobs_matrix, *params):
        Q = self.calc_exchangeability_matrix(word_probs, *params)
        Q *= mprobs_matrix
        row_totals = Q.sum(axis=1)
        Q -= numpy.diag(row_totals)
        Q *= 1.0 / (word_probs * row_totals).sum()
        return Q

class Empirical(StationaryQ, _ContinuousSubstitutionModel):
    """A continuous substitution model with a predefined instantaneous rate
    matrix."""

    @extend_docstring_from(_ContinuousSubstitutionModel.__init__)
    def __init__(self, alphabet, rate_matrix, **kw):
        """
        - rate_matrix: The instantaneous rate matrix
        """
        _ContinuousSubstitutionModel.__init__(self, alphabet, **kw)
        d = locals()
        exclude = ("self", "__class__")
        d = {k: v for k, v in d.items() if k not in exclude}
        self._serialisable.update(d)

        alphabet = self.get_alphabet()  # as may be altered by recode_gaps etc.
        N = len(alphabet)
        assert rate_matrix.shape == (N, N)
        assert numpy.alltrue(numpy.diagonal(rate_matrix) == 0)
        self._instantaneous_mask_f = rate_matrix * 1.0
        self._instantaneous_mask = self._instantaneous_mask_f != 0.0
        self.symmetric = _isSymmetrical(self._instantaneous_mask_f)
        self.parameter_order = []
        self.check_params_exist()

    def calc_exchangeability_matrix(self, mprobs):
        return self._instantaneous_mask_f.copy()

class Parametric(_ContinuousSubstitutionModel):
    """A continuous substitution model with only user-specified substitution
    parameters. This is a general process -- non-stationary and, if specified
    via predicates, non-reversible"""

    @extend_docstring_from(_ContinuousSubstitutionModel.__init__)
    def __init__(self, alphabet, predicates=None, scales=None, **kw):
        """
        - predicates: a dict of {name:predicate}. See cogent3.evolve.predicate
        - scales: scale rules, dict with predicates
        """
        self._canned_predicates = None
        _ContinuousSubstitutionModel.__init__(self, alphabet, **kw)

        d = locals()
        exclude = ("self", "__class__")
        d = {k: v for k, v in d.items() if k not in exclude}
        self._serialisable.update(d)

        (predicate_masks, predicate_order) = self._adapt_predicates(predicates or [])

        # Check for redundancy in predicates, ie: 1 or more than combine
        # to be equivalent to 1 or more others, or the distance params.
        # Give a clearer error in simple cases like always false or true.
        for (name, matrix) in list(predicate_masks.items()):
            if numpy.alltrue((matrix == 0).flat):
                raise ValueError("Predicate %s is always false." % name)
        predicates_plus_scale = predicate_masks.copy()
        predicates_plus_scale[None] = self._instantaneous_mask
        for (name, matrix) in list(predicate_masks.items()):
            if numpy.alltrue((matrix == self._instantaneous_mask).flat):
                raise ValueError("Predicate %s is always true." % name)
        if redundancy_in_predicate_masks(predicate_masks):
            raise ValueError("Redundancy in predicates.")
        if redundancy_in_predicate_masks(predicates_plus_scale):
            raise ValueError(
                "Some combination of predicates is"
                " equivalent to the overall rate parameter."
            )

        self.predicate_masks = predicate_masks
        self.parameter_order = []
        self.predicate_indices = []
        self.symmetric = _isSymmetrical(self._instantaneous_mask)
        for pred in predicate_order:
            mask = predicate_masks[pred]
            if not _isSymmetrical(mask):
                self.symmetric = False
            indices = numpy.nonzero(mask)
            assert numpy.alltrue(mask[indices] == 1)
            self.parameter_order.append(pred)
            self.predicate_indices.append(indices)
        (self.scale_masks, scale_order) = self._adapt_predicates(scales or [])
        self.check_params_exist()

    def calc_exchangeability_matrix(self, mprobs, *params):
        assert len(params) == len(self.predicate_indices), self.parameter_order
        R = self._instantaneous_mask_f.copy()
        for (indices, par) in zip(self.predicate_indices, params):
            R[indices] *= par
        return R

    def ascii_art(self, delim="", delim2="|", max_width=70, return_table=False):
        """An ASCII-art table representing the model.  'delim' delimits
        parameter names, 'delim2' delimits motifs"""
        from cogent3.util.table import Table

        labels = [m for m in self.alphabet]
        pars = self.get_matrix_params()
        rows = []
        for i, row in enumerate(pars):
            r = [labels[i]] + [delim.join(cell) for cell in row]
            r[i + 1] = "*"  # identity
            rows.append(r)

        labels.insert(0, r"From\To")
        if self.name:
            title = "%s rate matrix" % self.name
        else:
            title = "rate matrix"

        t = Table(
            header=labels,
            data=rows,
            max_width=max_width,
            title=title,
            index_name=r"From\To",
        )
        result = t if return_table else t.to_string(center=True)
        return result

    def get_matrix_params(self):
        """Return the parameter assignment matrix."""
        dim = len(self.alphabet)
        Pars = numpy.zeros([dim, dim], object)
        for x, y in [(x, y) for x in range(dim) for y in range(dim)]:
            Pars[x][y] = []  # a limitation of numpy.  [x,y] = [] fails!
            if not self._instantaneous_mask[x, y]:
                continue
            for par in self.predicate_masks:
                if self.predicate_masks[par][x, y]:
                    Pars[x, y].append(par)
            # sort the matrix entry to facilitate scaling calculations
            Pars[x, y].sort()
        return Pars

    def get_param_list(self):
        """Return a list of parameter names."""
        return list(self.predicate_masks.keys())

    def is_instantaneous(self, x, y):
        return self._is_instantaneous(x, y)

    def get_substitution_rate_value_from_Q(self, Q, motif_probs, pred):
        pred_mask = list(self._adapt_predicates([pred])[0].values())[0]
        pred_row_totals = numpy.sum(pred_mask * Q, axis=1)
        inst_row_totals = numpy.sum(self._instantaneous_mask * Q, axis=1)
        r = sum(pred_row_totals * motif_probs)
        t = sum(inst_row_totals * motif_probs)
        pred_size = numpy.sum(pred_mask.flat)
        inst_size = sum(self._instantaneous_mask.flat)
        return (r / pred_size) / ((t - r) / (inst_size - pred_size))

    def get_scaled_lengths_from_Q(self, Q, motif_probs, length):
        lengths = {}
        for rule in self.scale_masks:
            lengths[rule] = length * self.get_scale_from_Qs(
                [Q], [1.0], motif_probs, rule
            )
        return lengths

    def get_scale_from_Qs(self, Qs, bin_probs, motif_probss, rule):
        rule = self.get_predicate_mask(rule)
        weighted_scale = 0.0
        bin_probs = numpy.asarray(bin_probs)
        for (Q, bin_prob, motif_probs) in zip(Qs, bin_probs, motif_probss):
            row_totals = numpy.sum(rule * Q, axis=1)
            motif_probs = numpy.asarray(motif_probs)
            word_probs = self.calc_word_probs(motif_probs)
            scale = sum(row_totals * word_probs)
            weighted_scale += bin_prob * scale
        return weighted_scale

    def get_predefined_predicates(self):
        # overridden in subclasses
        return {"indel": predicate.parse("-/?")}

    def get_predefined_predicate(self, name):
        # Called by predicate parsing code
        if self._canned_predicates is None:
            self._canned_predicates = self.get_predefined_predicates()
        return self._canned_predicates[name].interpret(self)

    def _adapt_predicates(self, rules):
        # dict or list of callables, predicate objects or predicate strings
        if isinstance(rules, dict):
            rules = list(rules.items())
        else:
            rules = [(None, rule) for rule in rules]
        predicate_masks = {}
        order = []
        for (key, pred) in rules:
            (label, mask) = self.adapt_predicate(pred, key)
            if label in predicate_masks:
                raise KeyError('Duplicate predicate name "%s"' % label)
            predicate_masks[label] = mask
            order.append(label)
        return predicate_masks, order

    def adapt_predicate(self, pred, label=None):
        if isinstance(pred, str):
            pred = predicate.parse(pred)
        elif isinstance(pred, Callable):
            pred = predicate.UserPredicate(pred)
        pred_func = pred.make_model_predicate(self)
        label = label or repr(pred)
        mask = predicate2matrix(
            self.get_alphabet(), pred_func, mask=self._instantaneous_mask
        )
        return (label, mask)

    def get_predicate_mask(self, pred):
        if pred in self.scale_masks:
            mask = self.scale_masks[pred]
        elif pred in self.predicate_masks:
            mask = self.predicate_masks[pred]
        else:
            (label, mask) = self.adapt_predicate(pred)
        return mask

class Stationary(StationaryQ, Parametric):
    def __init__(self, *args, **kw):
        Parametric.__init__(self, *args, **kw)

class TimeReversible(Stationary):
    def __init__(self, *args, **kw):
        """"""
        Stationary.__init__(self, *args, **kw)
        if not self.symmetric:
            raise ValueError(
                "TimeReversible exchangeability terms must be fully balanced"
            )

class _TimeReversibleNucleotide(TimeReversible):
    def get_predefined_predicates(self):
        return {
            "transition": predicate.parse("R/R") | predicate.parse("Y/Y"),
            "transversion": predicate.parse("R/Y"),
            "indel": predicate.parse("-/?"),
            "kappa": (predicate.parse("R/R") | predicate.parse("Y/Y")).aliased("kappa"),
        }

class TimeReversibleNucleotide(_TimeReversibleNucleotide):
    """A nucleotide substitution model."""

    def __init__(self, *args, **kw):
        _TimeReversibleNucleotide.__init__(self, moltype.DNA.alphabet, *args, **kw)

class TimeReversibleDinucleotide(_TimeReversibleNucleotide):
    """A dinucleotide substitution model."""

    def __init__(self, *args, **kw):
        _TimeReversibleNucleotide.__init__(
            self, moltype.DNA.alphabet, motif_length=2, *args, **kw
        )

class TimeReversibleTrinucleotide(_TimeReversibleNucleotide):
    """A trinucleotide substitution model."""

    def __init__(self, *args, **kw):
        _TimeReversibleNucleotide.__init__(
            self, moltype.DNA.alphabet, motif_length=3, *args, **kw
        )

class TimeReversibleProtein(TimeReversible):
    """base protein substitution model."""

    def __init__(self, with_selenocysteine=False, *args, **kw):
        alph = moltype.PROTEIN.alphabet
        if not with_selenocysteine:
            alph = alph.get_subset("U", excluded=True)
        TimeReversible.__init__(self, alph, *args, **kw)

def EmpiricalProteinMatrix(
    matrix, motif_probs=None, optimise_motif_probs=False, recode_gaps=True, **kw
):
    alph = moltype.PROTEIN.alphabet.get_subset("U", excluded=True)
    return Empirical(
        alph,
        rate_matrix=matrix,
        motif_probs=motif_probs,
        model_gaps=False,
        recode_gaps=recode_gaps,
        optimise_motif_probs=optimise_motif_probs,
        **kw,
    )

class _CodonPredicates:
    """predicates for silent and replacement substitutions"""

    def __init__(self, gc):
        """
        Parameters
        ----------

        gc
            a genetic code instance
        """
        self.gc = gc

    def silent(self, x, y):
        return x != "---" and y != "---" and self.gc[x] == self.gc[y]

    def replacement(self, x, y):
        return x != "---" and y != "---" and self.gc[x] != self.gc[y]

class _Codon:
    long_indels_are_instantaneous = True

    def _is_instantaneous(self, x, y):
        if x == self.gapmotif or y == self.gapmotif:
            return x != y
        else:
            ndiffs = sum([X != Y for (X, Y) in zip(x, y)])
            return ndiffs == 1

    def get_predefined_predicates(self):
        codon_preds = _CodonPredicates(self.get_alphabet().get_genetic_code())

        preds = _TimeReversibleNucleotide.get_predefined_predicates(self)
        preds.update(
            {
                "indel": predicate.parse("???/---"),
                "silent": predicate.UserPredicate(codon_preds.silent),
                "replacement": predicate.UserPredicate(codon_preds.replacement),
                "omega": predicate.UserPredicate(codon_preds.replacement),
            }
        )
        return preds

class TimeReversibleCodon(_Codon, _TimeReversibleNucleotide):
    """Core substitution model for codons"""

    def __init__(self, alphabet=None, gc=None, **kw):
        if gc is not None:
            alphabet = moltype.CodonAlphabet(gc=gc)
        alphabet = alphabet or moltype.STANDARD_CODON
        _TimeReversibleNucleotide.__init__(self, alphabet, **kw)


# 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
import tqdm
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.autoguide.initialization import InitMessenger
from pyro.infer.reparam import LocScaleReparam
from pyro.nn.module import PyroModule, PyroParam
from pyro.ops.streaming import CountMeanVarianceStats, StatsOfDict
from pyro.optim import ClippedAdam
from pyro.poutine.util import site_is_subsample
from torch.distributions import constraints

import pyrocov.geo

from . import pangolin, sarscov2
from .ops import sparse_multinomial_likelihood
from .util import pearson_correlation, quotient_central_moments

# Requires https://github.com/pyro-ppl/pyro/pull/2953
try:
    from pyro.infer.autoguide.effect import AutoRegressiveMessenger
except ImportError:
    AutoRegressiveMessenger = object

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):
    """
    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, min_samples: int = 50) -> List[str]:
    """
    Compute a list of lineages ranked in descending order of cut size.
    This is used in growth rate leave-one-out prediction experiments.
    """

    def get_parent(lineage):
        lineage = pangolin.decompress(lineage)
        parent = pangolin.get_parent(lineage)
        if parent is not None:
            parent = pangolin.compress(parent)
        return parent

    # Compute sample counts.
    lineage_id_inv = full_dataset["lineage_id_inv"]
    lineage_id = full_dataset["lineage_id"]
    clade_counts = full_dataset["weekly_clades"].sum([0, 1])
    lineage_counts = clade_counts.new_zeros(len(lineage_id)).scatter_add_(
        0, full_dataset["clade_id_to_lineage_id"], clade_counts
    )
    weekly_clades = full_dataset["weekly_clades"]  # [T, P, C]
    lineage_counts = weekly_clades.sum([0, 1])  # [C]
    descendent_counts = lineage_counts.clone()
    for c, lineage in enumerate(lineage_id_inv):
        ancestor = get_parent(lineage)
        while ancestor is not None:
            a = lineage_id.get(ancestor)
            if a is not None:
                descendent_counts[a] += lineage_counts[c]
            ancestor = get_parent(ancestor)
    total = lineage_counts.sum().item()
    cut_size = torch.min(descendent_counts, total - descendent_counts)

    # Filter and sort lineages by cut size.
    ranked_lineages = [
        (size, lineage)
        for size, lineage in zip(cut_size.tolist(), lineage_id_inv)
        if lineage not in ("A", "B", "B.1")
        if size >= min_samples
    ]
    ranked_lineages.sort(reverse=True)
    return [name for gap, name in ranked_lineages]

def dense_to_sparse(x):
    index = x.nonzero(as_tuple=False).T.contiguous()
    value = x[tuple(index)]
    total = x.sum(-1)
    return {"index": index, "value": value, "total": total}

def load_gisaid_data(
    *,
    device="cpu",
    min_region_size=50,
    include={},
    exclude={},
    end_day=None,
    columns_filename="results/usher.columns.pkl",
    features_filename="results/usher.features.pt",
    feature_type="aa",
) -> dict:
    """
    Loads the two files columns_filename and features_filename,
    converts the input to PyTorch tensors and truncates the data according to
    ``include`` and ``exclude``.

    :param str device: torch device to use
    :param dict include: filters of data to include
    :param dict exclude: filters of data to exclude
    :param end_day: last day to include
    :param str columns_filename:
    :param str features_filename:
    :param str feature_type: Either "aa" for amino acid features or "nuc" for
        nucleotide features.
    :returns: A dataset dict
    :rtype: dict
    """
    logger.info("Loading data")
    include = include.copy()
    exclude = exclude.copy()

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

    # Load column data.
    with open(columns_filename, "rb") as f:
        columns = pickle.load(f)
    # Clean up location ids (temporary; this should be done in preprocess_gisaid.py).
    columns["location"] = list(map(pyrocov.geo.gisaid_normalize, columns["location"]))
    logger.info(f"Training on {len(columns['day'])} rows with columns:")
    logger.info(", ".join(columns.keys()))

    # Aggregate regions smaller than min_region_size to country level.
    fine_regions = get_fine_regions(columns, min_region_size)

    # Filter features into numbers of mutations and possibly genes.
    usher_features = torch.load(features_filename)
    mutations = usher_features[f"{feature_type}_mutations"]
    features = usher_features[f"{feature_type}_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))))

    # Construct the list of clades.
    clade_id_inv = usher_features["clades"]
    clade_id = {k: i for i, k in enumerate(clade_id_inv)}
    clades = columns["clade"]

    # Generate sparse_data.
    sparse_data: dict = Counter()
    countries = set()
    states = set()
    state_to_country_dict = {}
    location_id: dict = OrderedDict()
    skipped_clades = set()
    num_obs = 0
    for day, location, clade in zip(columns["day"], columns["location"], clades):
        if clade not in clade_id:
            if clade not in skipped_clades:
                skipped_clades.add(clade)
                if not clade.startswith("fine"):
                    logger.warning(f"WARNING skipping unsampled clade {clade}")
            continue

        # Filter by include/exclude
        row = {
            "location": location,
            "day": day,
            "clade": clade,
        }
        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)
        # Populate countries on the left and states on the right.
        if len(parts) == 2:  # country only
            countries.add(location)
            p = location_id.setdefault(location, len(countries) - 1)
        else:  # state and country
            country = " / ".join(parts[:2])
            countries.add(country)
            c = location_id.setdefault(country, len(countries) - 1)
            states.add(location)
            p = location_id.setdefault(location, -len(states))
            state_to_country_dict[p] = c

        # Save sparse data.
        num_obs += 1
        t = day // TIMESTEP
        c = clade_id[clade]
        sparse_data[t, p, c] += 1
    logger.warning(f"WARNING skipped {len(skipped_clades)} unsampled clades")
    state_to_country = torch.full((len(states),), 999999, dtype=torch.long)
    for s, c in state_to_country_dict.items():
        state_to_country[s] = c
    logger.info(f"Found {len(states)} states in {len(countries)} countries")
    location_id_inv = [None] * len(location_id)
    for k, i in location_id.items():
        location_id_inv[i] = k
    assert all(location_id_inv)

    # Generate weekly_clades 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)
    C = len(clade_id)
    weekly_clades = torch.zeros(T, P, C)
    for tps, n in sparse_data.items():
        weekly_clades[tps] = n
    logger.info(f"Dataset size [T x P x C] {T} x {P} x {C}")

    logger.info(
        f"Keeping {num_obs}/{len(clades)} rows "
        f"(dropped {len(clades) - int(num_obs)})"
    )

    # Construct sparse representation.
    pc_index = weekly_clades.ne(0).any(0).reshape(-1).nonzero(as_tuple=True)[0]
    sparse_counts = dense_to_sparse(weekly_clades)

    # Construct time scales centered around observations.
    time = torch.arange(float(T)) * TIMESTEP / GENERATION_TIME
    time -= time.mean()

    # Construct lineage <-> clade mappings.
    lineage_to_clade = usher_features["lineage_to_clade"]
    clade_to_lineage = usher_features["clade_to_lineage"]
    lineage_id_inv = sorted(lineage_to_clade)
    lineage_id = {k: i for i, k in enumerate(lineage_id_inv)}
    clade_id_to_lineage_id = torch.zeros(len(clade_to_lineage), dtype=torch.long)
    for c, l in clade_to_lineage.items():
        clade_id_to_lineage_id[clade_id[c]] = lineage_id[l]
    lineage_id_to_clade_id = torch.zeros(len(lineage_to_clade), dtype=torch.long)
    for l, c in lineage_to_clade.items():
        lineage_id_to_clade_id[lineage_id[l]] = clade_id[c]

    dataset = {
        "clade_id": clade_id,
        "clade_id_inv": clade_id_inv,
        "clade_id_to_lineage_id": clade_id_to_lineage_id,
        "clade_to_lineage": usher_features["clade_to_lineage"],
        "features": features,
        "lineage_id": lineage_id,
        "lineage_id_inv": lineage_id_inv,
        "lineage_id_to_clade_id": lineage_id_to_clade_id,
        "lineage_to_clade": usher_features["lineage_to_clade"],
        "location_id": location_id,
        "location_id_inv": location_id_inv,
        "mutations": mutations,
        "pc_index": pc_index,
        "sparse_counts": sparse_counts,
        "state_to_country": state_to_country,
        "time": time,
        "weekly_clades": weekly_clades,
    }
    return dataset

def subset_gisaid_data(
    gisaid_dataset: dict,
    location_queries=None,
    max_clades=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_clades"] = new["weekly_clades"].index_select(1, ids)

    # Select clades.
    if new["weekly_clades"].size(-1) > max_clades:
        ids = (
            new["weekly_clades"]
            .sum([0, 1])
            .sort(0, descending=True)
            .indices[:max_clades]
        )
        new["weekly_clades"] = new["weekly_clades"].index_select(-1, ids)
        new["features"] = new["features"].index_select(0, ids)
        new["clade_id_inv"] = [new["clade_id_inv"][i] for i in ids.tolist()]
        new["clade_id"] = {name: i for i, name in enumerate(new["clade_id_inv"])}
        new["sparse_counts"] = dense_to_sparse(new["weekly_clades"])

    # 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, {}/{} clades, {}/{} mutations, {}/{} samples".format(
            len(new["location_id"]),
            len(old["location_id"]),
            len(new["clade_id"]),
            len(old["clade_id"]),
            len(new["mutations"]),
            len(old["mutations"]),
            int(new["weekly_clades"].sum()),
            int(old["weekly_clades"].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_clades"])
    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 clade 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 clade
      portions in each (time, regin) bin.
    """
    # Tensor shapes are commented at at the end of some lines.
    features = dataset["features"]
    time = dataset["time"]  # [T]
    weekly_clades = dataset["weekly_clades"]
    sparse_counts = dataset["sparse_counts"]
    clade_id_to_lineage_id = dataset["clade_id_to_lineage_id"]
    pc_index = dataset["pc_index"]
    T, P, C = weekly_clades.shape
    C, F = features.shape
    L = len(dataset["lineage_id"])
    PC = len(pc_index)
    assert PC <= P * C
    assert time.shape == (T,)
    assert clade_id_to_lineage_id.shape == (C,)

    # Optionally extend time axis.
    if forecast_steps is not None:  # During prediction.
        T = T + forecast_steps
        t0 = time[0]
        dt = time[1] - time[0]
        time = t0 + dt * torch.arange(float(T))
        assert time.shape == (T,)

    clade_plate = pyro.plate("clade", C, dim=-1)
    place_plate = pyro.plate("place", P, dim=-2)
    time_plate = pyro.plate("time", T, dim=-3)
    pc_plate = pyro.plate("place_clade", PC, dim=-1)

    # Configure reparametrization (which does not affect model density).
    reparam = {}
    if "reparam" in model_type:
        reparam["coef"] = LocScaleReparam()
        reparam["pc_rate"] = LocScaleReparam()
        reparam["pc_init"] = LocScaleReparam()
    with poutine.reparam(config=reparam):

        # Sample global random variables.
        coef_scale = pyro.sample("coef_scale", dist.LogNormal(-4, 2))
        rate_scale = pyro.sample("rate_scale", dist.LogNormal(-4, 2))
        init_scale = pyro.sample("init_scale", dist.LogNormal(0, 2))

        # Assume relative growth rate depends strongly on mutations and weakly
        # on clade and place. Assume initial infections depend strongly on
        # clade and place.
        coef = pyro.sample(
            "coef", dist.Laplace(torch.zeros(F), coef_scale).to_event(1)
        )  # [F]
        with clade_plate:
            rate_loc = pyro.deterministic("rate_loc", 0.01 * coef @ features.T)  # [C]
        with pc_plate:
            pc_rate_loc = rate_loc.expand(P, C).reshape(-1)
            pc_rate = pyro.sample(
                "pc_rate", dist.Normal(pc_rate_loc[pc_index], rate_scale)
            )  # [PC]
            pc_init = pyro.sample("pc_init", dist.Normal(0, init_scale))  # [PC]
        with place_plate, clade_plate:
            rate = pyro.deterministic(
                "rate",
                pc_rate_loc.scatter(0, pc_index, pc_rate).reshape(P, C),
            )  # [P, C]
            init = pyro.deterministic(
                "init",
                torch.full((P * C,), -1e2).scatter(0, pc_index, pc_init).reshape(P, C),
            )  # [P, C]
        logits = init + rate * time[:, None, None]  # [T, P, C]

        # Optionally predict probabilities (during prediction).
        if forecast_steps is not None:
            probs = logits.new_zeros(logits.shape[:2] + (L,)).scatter_add_(
                -1, clade_id_to_lineage_id.expand_as(logits), logits.softmax(-1)
            )
            with time_plate, place_plate, pyro.plate("lineage", L, dim=-1):
                pyro.deterministic("probs", probs)
            return

        # Finally observe counts (during inference).
        if "dense" in model_type:  # equivalent either way
            # Compute a dense likelihood.
            with time_plate, place_plate:
                pyro.sample(
                    "obs",
                    dist.Multinomial(logits=logits.unsqueeze(-2), validate_args=False),
                    obs=weekly_clades.unsqueeze(-2),
                )  # [T, P, 1, C]
            return
        # Compromise between sparse and dense.
        logits = logits.log_softmax(-1)
        t, p, c = sparse_counts["index"]
        pyro.factor(
            "obs",
            sparse_multinomial_likelihood(
                sparse_counts["total"], logits[t, p, c], sparse_counts["value"]
            ),
        )

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_clades"].sum(0)  # [P, C]
        init.add_(1 / init.size(-1)).div_(init.sum(-1, True))
        init.log_().sub_(init.median(-1, True).values).add_(torch.randn(init.shape))
        self.init = init  # [P, C]
        self.init_decentered = init / 2
        self.init_loc = init.mean(0)  # [C]
        self.init_loc_decentered = self.init_loc / 2
        self.pc_init = self.init.reshape(-1)[dataset["pc_index"]]
        self.pc_init = self.pc_init / 2
        assert not torch.isnan(self.init).any()
        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_scale",
            "place_scale",
            "clade_scale",
        ):
            return torch.full(shape, 0.01)
        if name in (
            "rate_loc",
            "rate_loc_decentered",
            "coef",
            "coef_decentered",
            "rate",
            "rate_decentered",
            "pc_rate",
            "pc_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)
        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__(InitMessenger(init_loc_fn)(model))

        # Jointly estimate globals, mutation coefficients, and clade 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))

class RegressiveGuide(AutoRegressiveMessenger):
    def get_posterior(self, name, prior):
        if name == "coef":
            if not hasattr(self, "coef"):
                # Initialize.
                self.coef = PyroModule()
                n = prior.shape()[-1]
                rank = 100
                assert n > 1
                init_loc = self.init_loc_fn({"name": name, "fn": prior})
                self.coef.loc = PyroParam(init_loc, event_dim=1)
                self.coef.scale = PyroParam(
                    torch.full((n,), self._init_scale),
                    event_dim=1,
                    constraint=constraints.positive,
                )
                self.coef.cov_factor = PyroParam(
                    torch.empty(n, rank).normal_(0, 1 / rank ** 0.5),
                    event_dim=2,
                )
            scale = self.coef.scale
            cov_factor = self.coef.cov_factor * scale.unsqueeze(-1)
            cov_diag = scale * scale
            return dist.LowRankMultivariateNormal(self.coef.loc, cov_factor, cov_diag)

        return super().get_posterior(name, prior)

@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 not name.startswith("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 tqdm.tqdm(range(num_samples)):
            stats.update(get_conditionals(guide()))
        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)
    Elbo = JitTrace_ELBO if jit else Trace_ELBO
    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 == "regressive":
        guide = RegressiveGuide(model_, init_loc_fn=init_loc_fn, init_scale=0.01)
    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 or "prec_sqrts" 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 = Elbo(max_plate_nesting=3, ignore_jit_warnings=True)
    svi = SVI(model_, guide, optim, elbo)
    losses = []
    num_obs = dataset["weekly_clades"].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() < 1e8
    }
    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 clades.
    rate = quotient_central_moments(
        result["mean"]["rate"].mean(0), dataset["clade_id_to_lineage_id"]
    )[1]
    rate = rate - rate[dataset["lineage_id"]["A"]]
    for lineage in ["B.1.1.7", "B.1.617.2", "AY.23.1"]:
        R_RA = float(rate[dataset["lineage_id"][lineage]].exp())
        logger.info(f"R({lineage})/R(A) = {R_RA:0.3g}")
        stats[f"R({lineage})/R(A)"] = R_RA

    # Posterior predictive error.
    L = len(dataset["lineage_id"])
    weekly_clades = dataset["weekly_clades"]
    weekly_lineages = torch.zeros(weekly_clades.shape[:-1] + (L,)).scatter_add_(
        -1, dataset["clade_id_to_lineage_id"].expand_as(weekly_clades), weekly_clades
    )
    true = weekly_lineages + 1e-20  # avoid nans
    counts = true.sum(-1, True)
    true_probs = true / counts
    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"] = float(mae.sum(-1).mean())  # average over region
    stats["RMSE"] = float(mse.sum(-1).mean().sqrt())  # root average over region
    stats["KL"] = float(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, lineages 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"] = float(kl[p].sum() / true[:, p].sum())
        stats[f"{place} MAE"] = float(mae[p].sum())
        stats[f"{place} RMSE"] = float(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 lineage in lineages:
            i = dataset["lineage_id"][lineage]
            stats[f"{place} {lineage} MAE"] = mae[p, i]
            stats[f"{place} {lineage} RMSE"] = mse[p, i].sqrt()
            logger.info(
                "{} {}\tMAE = {:0.3g}, RMSE = {:0.3g}".format(
                    place,
                    lineage,
                    stats[f"{place} {lineage} MAE"],
                    stats[f"{place} {lineage} 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()}