#!/usr/local/bin/python

from pyrosetta import *
import re,sys
import os, shutil
import random
import numpy as np
import pickle
import math
os.environ["OPENBLAS_NUM_THREADS"] = "1"

#exit()
#datadir = sys.argv[1]
initial_model = sys.argv[1]

def main():

    init('-hb_cen_soft -relax:default_repeats 1 -default_max_cycles 50')

#    os.chdir(datadir)

    #print "here";

    scorefxn_fa=create_score_function('ref2015')
    scorefxn_fa.set_weight(rosetta.core.scoring.atom_pair_constraint, 5)
    scorefxn_fa.set_weight(rosetta.core.scoring.dihedral_constraint, 1)
    scorefxn_fa.set_weight(rosetta.core.scoring.angle_constraint, 1)

    mmap = MoveMap()
    mmap.set_bb(True)
    mmap.set_chi(True)
    mmap.set_jump(True)

    relax=rosetta.protocols.relax.FastRelax()
    relax.set_scorefxn(scorefxn_fa)
    relax.max_iter(200)
    relax.dualspace(True)
    relax.set_movemap(mmap)

    #refine without cst to estimate model accuracy
    pose = pose_from_file(initial_model + ".pdb")

    print("refine without cst")
    relax.apply(pose)
    sfile=initial_model+"_ref.sc"
    SS = open(sfile, "w")
    name=initial_model  + "_ref.pdb";
    inf=name + "\t" + str(scorefxn_fa(pose))
    SS.write(inf)
    SS.write("\n")
    pose.dump_pdb(name)
    SS.close()

#end of main

def select_model(filename):
    sc=[]
    decoy=[]
    i=0
    with open(filename, "r") as f:
        for line in f:
            line=line.rstrip()
            if(line[0] == "#"):
                continue
            else:
                b=re.split("\t", line)

                sc.append(float(b[1]))
                decoy.append(b[0])
                #print b[0], b[1]

    idx=np.argsort(sc)
    topN=5
    if(topN>len(idx)): topN=len(idx)
    final=[]
    k=1
    for i in idx:
        #print decoy[i]
        if(not os.path.isfile(decoy[i])): continue
        name="model"+str(k)+".pdb"
        shutil.copy(decoy[i], name)
        k +=1
        if(k>5): break

def fetch_subset(decoys, score, prob):
    nd=len(decoys)
    selected_decoys=[]
    selected_scores=[]
    for i in range(0, nd):
        d=decoys[i]
        Reg=re.compile(r'(.+).pdb')
        mo=Reg.search(d)
        #print mo.group(1)
        rst=re.split("_", mo.group(1))
        if(float(rst[1])==prob):
            selected_decoys.append(mo.group(1))
            selected_scores.append(score[i])

    idx=np.argsort(selected_scores)
    topN=10
    if(topN>len(idx)): topN=len(idx)
    final=[]
    for j in range(0,topN):
        i=idx[j]
        final.append(selected_decoys[i])
        #print selected_decoys[i], selected_scores[i]
    return final;

def output_data(pose, scorefxn, name, filename):
    inf=name + "\t" + str(scorefxn(pose))
    SS = open(filename, "a")
    SS.write(inf)
    SS.write("\n")
    SS.close()

    pose.dump_pdb(name)

def remove_clash(scorefxn, mover, pose):
    clash_score=float(scorefxn(pose))
    print("clash_score=", clash_score)
    if(clash_score>10):
        for nm in range(0, 5):
            mover.apply(pose)
            clash_score=float(scorefxn(pose))
            print("clash_score=", clash_score)
            if(clash_score<10): break

def run_min(cst_all, n_sets, pose, mover1, mover2):
    if(len(cst_all)==0):
        print("warning: empty constraint set")
        return

    random.shuffle(cst_all)
    b_size=int(len(cst_all)/n_sets)
    for i in range(0, len(cst_all), b_size):
        batch=cst_all[i:i+b_size]
        add_cst(pose, batch)
        mover1.apply(pose)
        mover2.apply(pose)

def fetch_cst(cst, nres, sep1, sep2, cut, flag):
    pcut=cut
    array=[]
    for line in cst:
        #print line
        line=line.rstrip()
        b=re.split("\s+", line)
        #print b[1],b[2],b[3],b[4]
        m=re.search('(?<=#).+', line)
        dcst=re.split("\s+", m.group(0))
        i=int(b[2])
        j=int(b[4])

        if(j==i):j=int(b[6]) #omega, phi
        if(j==i):j=int(b[8]) #theta
        if(b[0]=="Dihedral"): pcut=cut+0.5
        if(b[0]=="Angle"): pcut=cut+0.6
        if(b[0]=="AtomPair"): pcut=cut

        sep=abs(j-i)
        if(sep<sep1 or sep >=sep2): continue
        #print dcst[3]
        if(flag==1):
            if(float(dcst[3])>=pcut):
                array.append(line)
                #print line
        else:
            if(float(dcst[3])<pcut):
                array.append(line)

    return array

def compute_dist(pose, i, atmi, j, atmj):

    aai=pose.residue(i)
    xyz_i=aai.xyz(atmi)

    aaj=pose.residue(j)
    xyz_j=aaj.xyz(atmj)
    d=dist(xyz_i, xyz_j)
    return d

def get_pairwise_dist(pose):
    nres=pose.total_residue()
    for i in range(1, nres-2):
        aai=pose.residue(i)
        resi=aai.name()
        resi=resi[0:3]
        atmi="CB"
        if(resi == "GLY"):atmi="CA"
        xyz_i=aai.xyz(atmi)
        for j in range(i+3, nres+1):
            aaj=pose.residue(j)
            resj=pose.residue(j).name()
            resj=resj[0:3]
            atmj="CB"
            if(resj == "GLY"):atmj="CA"
            xyz_j=aaj.xyz(atmj)
            d=dist(xyz_i, xyz_j)
            if(d<20):
                pass
                #print i, j, d
        #end j
    #end i

def dist(x, y):
    cut=20
    d=100
    a=abs(x[0]-y[0]);
    b=abs(x[1]-y[1]);
    c=abs(x[2]-y[2]);
    if(a>cut or b>cut or c>cut):
        return d
    else:
        d=math.sqrt(a**2+b**2+c**2)
        return d

def apply_cst(cst_file, b_size, pose, mover1, mover2):
    cst_all=read_cst(cst_file)
    if(b_size==0):
        b_size=len(cst_all)
    else:
        random.shuffle(cst_all)

    for i in range(0, len(cst_all), b_size):
        batch=cst_all[i:i+b_size]
        add_cst(pose, batch)
        #print(pose.constraint_set())
        mover1.apply(pose)
        mover2.apply(pose)
        #print(mover.num_accepts(), mover.acceptance_rate())

    return cst_all

def make_start(k, sequence, scorefxn, mover):

    pose=pose_from_sequence(sequence, 'centroid' )
    if(os.path.isfile("tor_ss.dat") and k<1):
        set_predicted_dihedral("tor_ss.dat", pose)
    else:
        set_random_dihedral(pose)

    clash_score=float(scorefxn(pose))
    if(clash_score>10):
        for nm in range(0, 5):
            mover.apply(pose)
            clash_score=float(scorefxn(pose))
            print("clash_score=", clash_score)
            if(clash_score<10): break

    return pose

def set_predicted_dihedral(filename,pose):
    with open(filename, "r") as f:
        for line in f:
            if(line[0] == "#"):
                continue
            else:
                b=re.split("\t", line)
                pose.set_phi(int(b[0]),float(b[2]))
                pose.set_psi(int(b[0]),float(b[3]))
                if(b[1]=="C"):b[1]="L"
                pose.set_secstruct(int(b[0]), b[1])
    return(pose)

def convert_spx(filename):
    out=open("tor_ss.dat", "w")
    out.write("#ID\tSS\tPHI\tPSI\n")
    with open(filename, "r") as f:
        for line in f:
            if(line[0] == "#"):
                continue
            else:
                b=re.split("\s+", line)

                out.write("%d\t%c\t%.1f\t%.1f\n" %(int(b[1]), b[3], float(b[4]), float(b[5])))

    out.close()

def set_random_dihedral(pose):
    nres = pose.total_residue()
    for i in range(1, nres):
	#pick phi/psi randomly from:
        #-140  153 180 0.135 B
	# -72  145 180 0.155 B
        #-122  117 180 0.073 B
	# -82  -14 180 0.122 A
	# -61  -41 180 0.497 A
	#  57   39 180 0.018 L
        phi,psi=random_dihedral()
        pose.set_phi(i,phi)
        pose.set_psi(i,psi)
        pose.set_omega(i,180)

    return(pose)

def random_dihedral():
    phi=0
    psi=0
    r=random.random()
    if(r<=0.135):
        phi=-140
        psi=153
    elif(r>0.135 and r<=0.29):
        phi=-72
        psi=145
    elif(r>0.29 and r<=0.363):
        phi=-122
        psi=117
    elif(r>0.363 and r<=0.485):
        phi=-82
        psi=-14
    elif(r>0.485 and r<=0.982):
        phi=-61
        psi=-41
    else:
        phi=57
        psi=39
    return(phi, psi)

def read_cst(file):
    array=[]
    with open(file, "r") as f:
        for line in f:
            #print line
            line=line.rstrip()
            array.append(line)
    return array

def add_cst(pose, array):
    constraints = rosetta.protocols.constraint_movers.ConstraintSetMover()
    tmpname=initial_model+"_"+"tmp.sp";
    F = open(tmpname, "w")
    for a in array:
        F.write(a)
        F.write("\n")
    F.close()
    constraints.constraint_file(tmpname)
    constraints.add_constraints(True)
    constraints.apply(pose)
    os.remove(tmpname)

def read_fasta(file):
    fasta="";
    with open(file, "r") as f:
        for line in f:
            if(line[0] == ">"):
                continue
            else:
                line=line.rstrip()
                fasta = fasta + line;
    return fasta

if __name__ == '__main__':
    main()


from geo.Mesh import Mesh
import os
import torch
import torch.nn.functional as F
import imageio
import numpy as np

class FaceModel:
    '''
    FaceModel is the class bridging between torch-tensor uv representation of face,
    with traditional geometry representation for rendering pipelines.
    FaceModel stores all its elements using numpy array,
    and provides functions converting from tensors.
    '''

    def __init__(self, template_file=os.path.join(os.path.dirname(__file__), 'template', 'template.obj')):
        self.mesh = Mesh(path=template_file)

        # remove backside of head
        frontal_vt_flag = self.mesh.tex_coords[:, 0] < 1
        frontal_face_flag = frontal_vt_flag[self.mesh.tex_faces].sum(axis=1) == 3
        self.mesh.tex_faces = self.mesh.tex_faces[frontal_face_flag]
        self.mesh.faces = self.mesh.faces[frontal_face_flag]

        # now deal with albedo
        template_uv_file = os.path.join(os.path.dirname(__file__), 'template', 'uv.png')
        self.albedo = imageio.imread(template_uv_file)

    def update_pc(self, pc):
        '''
        Update the mesh geometry by providing pc tensor
        :param pc: [3, H, W] detached tensor of pc
        '''
        uv = self.mesh.tex_coords * 2 - 1
        # This is important since the image array has (0, 0) index on top-left corner of the image
        uv[:, 1] *= -1
        uv = torch.from_numpy(uv).float()
        uv = uv.unsqueeze(0).unsqueeze(0)
        tmp = F.grid_sample(pc.unsqueeze(0), uv)
        uv_pos = tmp[0, :, 0, :].transpose(0, 1).numpy()
        self.mesh.vertices = uv_pos
        self.mesh.faces = self.mesh.tex_faces
        self.mesh.re_calculate_normal()

    def update_al(self, al):
        '''
        Update the albedo by providing al tensor
        :param al: [3, H, W] detached tensor of al, re-normalized to [0, 1]
        '''

        def to_numpy(tensor):
            return tensor.detach().cpu().numpy().transpose(1, 2, 0)

        self.albedo = to_numpy(al)

    def export_to_obj(self, path):
        mtl_path = path.replace('.obj', '.mtl')
        al_path = path.replace('.obj', '.albedo.png')

        # write obj
        file = open(path, 'w')
        file.write('mtllib ./%s\n' % os.path.basename(mtl_path))
        for v in self.mesh.vertices:
            file.write('v %.6f %.6f %.6f\n' % (v[0], v[1], v[2]))
        for vt in self.mesh.tex_coords:
            file.write('vt %.6f %.6f\n' % (vt[0], vt[1]))

        for f, ft in zip(self.mesh.faces, self.mesh.tex_faces):
            f_plus = f + 1
            ft_plus = ft + 1
            file.write('f %d/%d %d/%d %d/%d\n' % (f_plus[0], ft_plus[0], f_plus[1], ft_plus[1], f_plus[2], ft_plus[2]))
        file.close()

        # write mtl
        file = open(mtl_path, 'w')
        file.write('''newmtl Textured
Ka 1.000 1.000 1.000
Kd 1.000 1.000 1.000
Ks 0.000 0.000 0.000
d 1.0
illum 2
''')
        file.write('map_Ka %s\n' % os.path.basename(al_path))
        file.close()

        imageio.imwrite(al_path, self.albedo)

class FaceVis:
    gl_content = False
    resolution = 0

    def __init__(self, resolution=512):
        if not self.gl_content:
            from render.gl import create_opengl_context
            create_opengl_context(resolution, resolution)
            self.gl_content = True
            self.resolution = 512
        from render.gl import ShRender
        self.render = ShRender(self.resolution, self.resolution)

    def render_model(self, face_model):
        # Init a camera
        from render.BaseCamera import BaseCamera
        from render.CameraPose import CameraPose

        intrinsic_cam = BaseCamera()
        intrinsic_cam.set_parameters(40)
        intrinsic_cam.near = 20
        intrinsic_cam.far = 60
        extrinsic_cam = CameraPose()
        extrinsic_cam.center = [0, 0, 40]
        extrinsic_cam.sanity_check()

        vertices = face_model.mesh.vertices
        faces = face_model.mesh.faces
        normals = face_model.mesh.normals
        norm_faces = face_model.mesh.norm_faces
        tex_coords = face_model.mesh.tex_coords
        tex_faces = face_model.mesh.tex_faces

        # Here we pack the vertex data needed for the render
        vert_data = vertices[faces.reshape([-1])]
        self.render.set_attrib(0, vert_data)
        norm_data = normals[norm_faces.reshape([-1])]
        self.render.set_attrib(1, norm_data)
        uv_data = tex_coords[tex_faces.reshape([-1])]
        self.render.set_attrib(2, uv_data)

        # Here we set the textures
        self.render.set_texture('TargetTexture', face_model.albedo)

        # Use a preset sh
        sh = np.array([[0.865754, 0.880196, 0.947154],
                       [-0.205633, -0.211215, -0.250504],
                       [0.349584, 0.365084, 0.463253],
                       [-0.0789622, -0.0750734, -0.091405],
                       [0.077691, 0.0810771, 0.0922284],
                       [-0.402685, -0.40834, -0.462763],
                       [0.328907, 0.328656, 0.370725],
                       [-0.131815, -0.140992, -0.158298],
                       [-0.0992293, -0.0983686, -0.107975]])

        # Here we pack all uniform variables of the shader into a single dictionary
        model_view = extrinsic_cam.get_model_view_mat()
        persp = intrinsic_cam.get_projection_mat()
        uniform_dict = {
            'ModelMat': model_view,
            'PerspMat': persp,
            'SHCoeffs': sh,
        }
        self.render.draw(
            uniform_dict
        )

        # Finally we save the rendered results.
        # Compare with the standard output under data/tests/
        color = self.render.get_color()
        return color


# Copyright (c) 2017, Carnegie Mellon University. All rights reserved.
#
# Use of the K-NRM package is subject to the terms of the software license set
# forth in the LICENSE file included with this software, and also available at
# https://github.com/AdeDZY/K-NRM/blob/master/LICENSE

"""
ranking data generator
include:
    point wise generator, each time yield a:
        query-doc pair (x), y
        x is a dict: h['q'] = array of q (batch size * max q len)
                    h['d'] = array of d (batch size * max title len)
                    h['idf] = None or array of query term idf (batch size * max q len)

    pair wise generator, each time yield a
        query-doc+, doc- pair (x), y
            the same x with pointwise's, with an additional key:
                h['d_aux'] = array of the second d (batch size * max title len)

    all terms are int ids

"""

from traitlets.config import Configurable
from traitlets import (
    Unicode,
    Int,
    List,
    Float,
    Bool,
)
import numpy as np
from numpy import genfromtxt
import logging
from io import StringIO
import sys
reload(sys)
sys.setdefaultencoding('UTF8')

class DataGenerator(Configurable):
    title_in = Unicode('/bos/data1/sogou16/data/training/1m_title.pad_t50',
                       help='titles term id csv, must be padded').tag(config=True)
    max_q_len = Int(10, help='max q len').tag(config=True)
    max_d_len = Int(50, help='max document len').tag(config=True)
    q_name = Unicode('q')
    d_name = Unicode('d')
    aux_d_name = Unicode('d_aux')
    idf_name = Unicode('idf')
    neg_sample = Int(1, help='negative sample').tag(config=True)
    load_litle_pool=Bool(False, help='load little pool at beginning').tag(conf=True)
    min_score_diff = Float(0, help='min score difference for click data generated pairs').tag(config=True)
    vocabulary_size = Int(2000000).tag(config=True)

    def __init__(self, **kwargs):
        super(DataGenerator, self).__init__(**kwargs)
        self.m_title_pool = np.array(None)
        if self.load_litle_pool and self.neg_sample:
            self._load_title_pool()
        print "min_score_diff: ", self.min_score_diff
        print "generator's vocabulary size: ", self.vocabulary_size

    def _load_title_pool(self):
        if self.title_in:
            logging.info('start loading title pool [%s]', self.title_in)
            self.m_title_pool = genfromtxt(self.title_in, delimiter=',',  dtype=int,)
            logging.info('loaded [%d] title pool', self.m_title_pool.shape[0])

    def pointwise_generate(self, pair_stream, batch_size, with_label=True, with_idf=False):
        """
        to use: initial the generator = ClickDataGenerator(config=conf)
            and then for X,Y in generator.pointwise_generator(pair_stream, batch_size)
        :param pair_stream: the (probably infinite) stream of query \t clicked url
            e.g. itertools.cycle(open(file))
        :param batch_size: int, a batch size
        :param with_label: if True, then there is a third column in pair_stream: \t label (int)
        :param with_idf: if True, the third col in pair_stream is the query term idf
        :return: yield a batched X and Y
        """
        l_q = []
        l_d = []
        l_idf = []
        l_y = []
        for line in pair_stream:
            cols = line.split('\t')
            q = np.array([int(t) for t in cols[0].split(',')])  #np.loadtxt(StringIO(unicode(cols[0])), delimiter=',',  dtype=int,)
            doc = np.array([int(t) for t in cols[1].split(',')])  #np.loadtxt(StringIO(unicode(cols[1])), delimiter=',',  dtype=int,)

            if with_idf:
                idf = np.array([int(t) for t in cols[2].split(',')])
            y = 0
            if with_label:
                y = int(cols[-1])
            v_q = np.ones(self.max_q_len) * -1
            v_d = np.ones(self.max_d_len) * -1
            v_q[:min(q.shape[0], self.max_q_len)] = q[0:min(q.shape[0], self.max_q_len)]
            v_d[:min(doc.shape[0], self.max_d_len)] = doc[0:min(doc.shape[0], self.max_d_len)]

            l_q.append(v_q)
            l_d.append(v_d)
            l_y.append(y)

            if with_idf:
                v_idf = np.zeros(self.max_q_len)
                v_idf[:idf.shape[0]] = idf[0:min(q.shape[0], self.max_q_len)]
                l_idf.append(v_idf)

            if len(l_q) >= batch_size:
                Q = np.array(l_q,  dtype=int,)
                D = np.array(l_d,  dtype=int,)
                IDF = None
                if with_idf:
                    IDF = np.array(l_idf, dtype=float)
                Y = np.array(l_y,  dtype=int,)
                X = {self.q_name: Q, self.d_name: D, self.idf_name: IDF}
                yield X, Y
                l_q, l_d, l_y, l_idf = [], [], [], []
        if l_q:
            Q = np.array(l_q,  dtype=int,)
            D = np.array(l_d,  dtype=int,)
            IDF = None
            if with_idf:
                IDF = np.array(l_idf, dtype=float)
            Y = np.array(l_y,  dtype=int,)
            X = {self.q_name: Q, self.d_name: D, self.idf_name: IDF}
            yield X, Y
        logging.info('point wise generator to an end')

    def pairwise_generate(self, pair_stream, batch_size, with_idf=False):
        """
        to use: initial the generator = ClickDataGenerator(config=conf)
            and then for X,Y in generator.pairwise_generate(pair_stream, batch_size)
        :param pair_stream: the (probably infinite) stream of query \t clicked url
            e.g. itertools.cycle(open(file))
        :param batch_size: must be neg_sample * k size
        :param with_idf: if True, the third col in pair_stream is the query term idf
        :return: yield a batched X and Y
            NOTE: the Y is always 1, the order of pos and neg docs are not shuffled yet.
        """

        assert batch_size % self.neg_sample == 0
        pos_batch_size = batch_size / self.neg_sample

        for pos_X, pos_Y in self.pointwise_generate(pair_stream, pos_batch_size, with_label=False, with_idf=with_idf):
            idx = np.random.randint(self.m_title_pool.shape[0], size=batch_size)
            aux_D = self.m_title_pool[idx, :]
            new_Q = np.repeat(pos_X[self.q_name], self.neg_sample, axis=0)
            new_D = np.repeat(pos_X[self.d_name], self.neg_sample, axis=0)

            new_IDF = None
            if with_idf:
                new_IDF = np.repeat(pos_X[self.idf_name], self.neg_sample, axis=0)

            X = {self.q_name: new_Q, self.d_name: new_D, self.aux_d_name: aux_D, self.idf_name:new_IDF}
            Y = np.ones(batch_size)
            yield X, Y

    def pairwise_reader(self, pair_stream, batch_size, with_idf=False):
        l_q = []
        l_d = []
        l_d_aux = []
        l_idf = []
        l_y = []
        for line in pair_stream:
            cols = line.strip().split('\t')
            if len(cols) < 4: continue
            flag = True
            for col in cols:
                if not col.strip():
                    flag = False
                    break
            if not flag:
                print line
                continue
            y = float(cols[3])
            if abs(y) < self.min_score_diff:
                continue
            q = np.array([int(t) for t in cols[0].split(',') if int(t) < self.vocabulary_size])
            t1 = np.array([int(t) for t in cols[1].split(',') if int(t) < self.vocabulary_size])
            t2 = np.array([int(t) for t in cols[2].split(',') if int(t) < self.vocabulary_size])
            if y > 0:
                y = 1
            else:
                t1, t2 = t2, t1  # make the first always positive
                y = 1
            if with_idf:
                if len(cols) < 5:
                    idf = np.ones(len(q))
                else:
                    idf = np.array([float(t) for t in cols[4].split(',')])
            v_q = np.ones(self.max_q_len) * -1
            v_d = np.ones(self.max_d_len) * -1
            v_d_aux = np.ones(self.max_d_len) * -1
            v_q[:min(q.shape[0], self.max_q_len)] = q[:min(q.shape[0], self.max_q_len)]
            v_d[:min(t1.shape[0], self.max_d_len)] = t1[:min(t1.shape[0], self.max_d_len)]
            v_d_aux[:min(t2.shape[0], self.max_d_len)] = t2[:min(t2.shape[0], self.max_d_len)]

            l_q.append(v_q)
            l_d.append(v_d)
            l_d_aux.append(v_d_aux)
            l_y.append(y)

            if with_idf:
                v_idf = np.zeros(self.max_q_len)
                v_idf[:min(idf.shape[0], self.max_q_len)] = idf[:min(idf.shape[0], self.max_q_len)]
                l_idf.append(v_idf)

            if len(l_q) >= batch_size:
                Q = np.array(l_q,  dtype=int,)
                D = np.array(l_d,  dtype=int,)
                D_aux = np.array(l_d_aux, dtype=int)
                IDF = None
                if with_idf:
                    IDF = np.array(l_idf, dtype=float)
                Y = np.array(l_y,  dtype=int,)
                X = {self.q_name: Q, self.d_name: D, self.idf_name: IDF, self.aux_d_name: D_aux}
                yield X, Y
                l_q, l_d, l_d_aux, l_y, l_idf = [], [], [], [], []
        if l_q:
            Q = np.array(l_q,  dtype=int,)
            D = np.array(l_d,  dtype=int,)
            D_aux = np.array(l_d_aux,  dtype=int,)
            IDF = None
            if with_idf:
                IDF = np.array(l_idf, dtype=float)
            Y = np.array(l_y,  dtype=int,)
            X = {self.q_name: Q, self.d_name: D, self.idf_name: IDF, self.aux_d_name: D_aux}
            yield X, Y
        logging.info('pair wise reader to an end')

if __name__ == '__main__':
    from deeplearning4ir.utils import set_basic_log, load_py_config
    set_basic_log()
    if 4 != len(sys.argv):
        print "I test generator"
        print "3 para: config + click pair with int term + batch_size"
        DataGenerator.class_print_help()
        sys.exit(-1)
    conf = load_py_config(sys.argv[1])
    generator = DataGenerator(config=conf)

    pair_stream = open(sys.argv[2])
    batch_size = int(sys.argv[3])
    X, Y = next(generator.pointwise_generate(pair_stream, batch_size, False))
    a = np.ones(1)

    print "point wise Y:\n %s" % (np.array2string(Y))
    print "\n\n"
    print "q: \n %s" % (np.array2string(X[generator.q_name]))
    print "\n\n"
    print "d: \n %s" % (np.array2string(X[generator.d_name]))
    print "\n\n"
    X, Y = next(generator.pairwise_generate(pair_stream, batch_size))

    print "pair wise Y:\n %s" % (np.array2string(Y))
    print "\n\n"
    print "q: \n %s" % (np.array2string(X[generator.q_name]))
    print "\n\n"
    print "d: \n %s" % (np.array2string(X[generator.d_name]))
    print "\n\n"
    print "aux d: \n %s" % (np.array2string(X[generator.aux_d_name]))


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
convert ghostnet pytorch pretrained weights to corresponding tf.keras model weights.
tf.keras model definition should be prepared and the layer/op name should be matched
with pytorch state dict keys.
"""
import os, sys, argparse
import numpy as np
import torch
from ghostnet_pytorch import ghostnet

sys.path.append(os.path.join(os.path.dirname(os.path.realpath(__file__)), '..', '..', '..'))
from common.backbones.ghostnet import GhostNet

def ghostnet_pytorch_to_keras(output_file):
    # load pretrained pytorch model
    pytorch_model = ghostnet(pretrained=True, weights_path=None)
    pytorch_model_dict = {}

    # parse valid pytorch model weights and store in dict
    for k, v in pytorch_model.state_dict().items():
        if len(v.shape) > 0:
            pytorch_model_dict[k] = v
            #print(k, v.shape)

    # create empty keras model
    model = GhostNet(include_top=True, input_shape=(224, 224, 3), width=1.0, weights=None)

    # walk through keras model layers and get corresponding pytorch weights
    for layer in model.layers:
        lname = layer.name
        weights = np.asarray(model.get_layer(lname).get_weights())
        if len(weights) > 0:
            print('Parsing layer {} with shape {}'.format(lname, weights.shape))

            # adjust layer name to match pytorch weights key
            pytorch_lname = lname.replace('_', '.')
            pytorch_lname = pytorch_lname.replace('conv.stem', 'conv_stem')
            pytorch_lname = pytorch_lname.replace('primary.conv', 'primary_conv')
            pytorch_lname = pytorch_lname.replace('cheap.operation', 'cheap_operation')
            pytorch_lname = pytorch_lname.replace('conv.dw', 'conv_dw')
            pytorch_lname = pytorch_lname.replace('bn.dw', 'bn_dw')
            pytorch_lname = pytorch_lname.replace('conv.reduce', 'conv_reduce')
            pytorch_lname = pytorch_lname.replace('conv.expand', 'conv_expand')
            pytorch_lname = pytorch_lname.replace('conv.head', 'conv_head')

            # convert normal conv without bias
            if pytorch_lname.endswith(('conv', 'conv_stem', 'primary_conv.0', 'shortcut.2')):
                weight = pytorch_model_dict[pytorch_lname + '.weight'].cpu().numpy()
                weight = weight.transpose((2, 3, 1, 0))
                model.get_layer(lname).set_weights([weight])

            # convert depthwise conv without bias
            elif pytorch_lname.endswith(('cheap_operation.0', 'conv_dw', 'shortcut.0')):
                weight = pytorch_model_dict[pytorch_lname + '.weight'].cpu().numpy()
                weight = weight.transpose((2, 3, 0, 1))
                model.get_layer(lname).set_weights([weight])

            # convert normal conv with bias
            elif pytorch_lname.endswith(('conv_reduce', 'conv_expand', 'conv_head')):
                weight = pytorch_model_dict[pytorch_lname + '.weight'].cpu().numpy()
                weight = weight.transpose((2, 3, 1, 0))
                bias = pytorch_model_dict[pytorch_lname + '.bias'].cpu().numpy()
                model.get_layer(lname).set_weights([weight, bias])

            # convert batchnorm
            elif pytorch_lname.endswith(('bn1', 'primary_conv.1', 'cheap_operation.1', 'bn_dw', 'shortcut.1', 'shortcut.3')):
                weight = pytorch_model_dict[pytorch_lname + '.weight'].cpu().numpy()
                bias = pytorch_model_dict[pytorch_lname + '.bias'].cpu().numpy()
                running_mean = pytorch_model_dict[pytorch_lname + '.running_mean'].cpu().numpy()
                running_var = pytorch_model_dict[pytorch_lname + '.running_var'].cpu().numpy()
                model.get_layer(lname).set_weights([weight, bias, running_mean, running_var])

            # convert batchnorm
            elif pytorch_lname.endswith('classifier'):
                weight = pytorch_model_dict[pytorch_lname + '.weight'].cpu().numpy()
                bias = pytorch_model_dict[pytorch_lname + '.bias'].cpu().numpy()
                weight = weight.transpose((1, 0))
                model.get_layer(lname).set_weights([weight, bias])

            # unsupported layer
            else:
                raise ValueError('Unsupported layer name', pytorch_lname)

    model.summary()
    model.save_weights(output_file)
    print('Done')
    return

def main():
    parser = argparse.ArgumentParser(description='convert ghostnet pytorch pretrained weights to tf.keras weights')
    parser.add_argument('--output_file', required=True, type=str, help='output tf.keras weights file')
    args = parser.parse_args()

    ghostnet_pytorch_to_keras(args.output_file)

if __name__ == "__main__":
    main()


#! /usr/bin/env python3

"""
Prediction part of the algorithm.
"""

import tensorflow as tf
import numpy as np
import os
import preprocessing as pp
import analysis as an
from tensorflow.contrib import learn
import csv
from sklearn import metrics
import pandas as pd
import matplotlib.pyplot as plt
from pandas_ml import ConfusionMatrix
import yaml
import logging

# Constants
# ============================================

SEMEVAL_FOLDER = '../data/SemEval/Subtask1'
RESTAURANT_TRAIN = os.path.join(SEMEVAL_FOLDER, 'restaurant', 'train.xml')
RESTAURANT_TEST = os.path.join(SEMEVAL_FOLDER, 'restaurant', 'test',
                               'test_gold.xml')
LAPTOP_TRAIN = os.path.join(SEMEVAL_FOLDER, 'laptop', 'train.xml')
LAPTOP_TEST = os.path.join(SEMEVAL_FOLDER, 'laptop', 'test', 'test_gold.xml')
RESTAURANT_ENTITIES = ['FOOD', 'DRINKS', 'SERVICE', 'RESTAURANT', 'AMBIENCE',
                       'LOCATION']
LAPTOP_ENTITIES = ['LAPTOP', 'HARDWARE', 'SHIPPING', 'COMPANY', 'SUPPORT',
                   'SOFTWARE']
POLARITY = ['positive', 'neutral', 'negative']
RESTAURANT_ASPECTS = [
        'RESTAURANT#GENERAL', 'RESTAURANT#PRICES', 'RESTAURANT#MISCELLANEOUS',
        'FOOD#PRICES', 'FOOD#QUALITY', 'FOOD#STYLE_OPTIONS',
        'DRINKS#PRICES', 'DRINKS#QUALITY', 'DRINKS#STYLE_OPTIONS',
        'AMBIENCE#GENERAL', 'SERVICE#GENERAL', 'LOCATION#GENERAL']

# Functions
# ==================================================

def softmax(x):
    """Compute softmax values for each sets of scores in x."""
    if x.ndim == 1:
        x = x.reshape((1, -1))
    max_x = np.max(x, axis=1).reshape((-1, 1))
    exp_x = np.exp(x - max_x)
    return exp_x / np.sum(exp_x, axis=1).reshape((-1, 1))

def prediction_process_CNN(folderpath_run, config_file, focus):
    """
    Process predictions for one CNN in order to obtain some measures about
    its efficiency.
    :param folderpath_run: The filepath of a run of train.py.
    :param config_file: The configuration file of the project opened with yaml
    library.
    :param focus: (required) 'feature' or 'polarity'. This precises the
    folder of a CNN. It will lead to the folder 'CNN_feature' or
    'CNN_polarity'.
    :type focus: string
    :return: datasets['data'], all_predictions, datasets['target_names'].
    datasets['data'] are the sentences before cleaning (after cleaning it is
    x_raw), all_predictions represents the prediction of the algorithm
    depending on the focus and datasets['target_names'] are the labels
    possible for the predictions.
    """

    datasets = None

    # Load data
    dataset_name = config_file["datasets"]["default"]
    if dataset_name == "semeval":
        current_domain = config_file["datasets"][dataset_name]["current_domain"]
        if current_domain == 'RESTAURANT':
            datasets = pp.get_dataset_semeval(RESTAURANT_TEST, focus,
                                              FLAGS.aspects)
        elif current_domain == 'LAPTOP':
            datasets = pp.get_dataset_semeval(LAPTOP_TEST, focus)
        else:
            raise ValueError("The 'current_domain' parameter in the " +
                             "'config.yml' file must be 'RESTAURANT' " +
                             "or 'LAPTOP'")

    x_raw, y_test = pp.load_data_and_labels(datasets)
    y_test = np.argmax(y_test, axis=1)
    logger.debug("Total number of test examples: {}".format(len(y_test)))

    # Map data into vocabulary
    vocab_path = os.path.join(folderpath_run, 'CNN_' + focus, 'vocab')
    vocab_processor = learn.preprocessing.VocabularyProcessor.restore(
            vocab_path)
    x_test = np.array(list(vocab_processor.transform(x_raw)))

    logger.info("")
    logger.info("Evaluation :")
    logger.info("")

    # Evaluation
    # ==================================================
    checkpoints_folder = os.path.join(folderpath_run, 'CNN_' + focus,
                                      'checkpoints')
    checkpoint_file = tf.train.latest_checkpoint(checkpoints_folder)
    graph = tf.Graph()

    with graph.as_default():

        session_conf = tf.ConfigProto(
          allow_soft_placement=FLAGS.allow_soft_placement,
          log_device_placement=FLAGS.log_device_placement)
        sess = tf.Session(config=session_conf)

        with sess.as_default():

            # Load the saved meta graph and restore variables
            saver = tf.train.import_meta_graph(
                    "{}.meta".format(checkpoint_file))
            saver.restore(sess, checkpoint_file)

            # Get the placeholders from the graph by name
            input_x = graph.get_operation_by_name("input_x").outputs[0]
            # input_y = graph.get_operation_by_name("input_y").outputs[0]
            dropout_keep_prob = graph.get_operation_by_name(
                    "dropout_keep_prob").outputs[0]

            # Tensors we want to evaluate
            scores = graph.get_operation_by_name("output/scores").outputs[0]

            # Tensors we want to evaluate
            predictions = graph.get_operation_by_name(
                    "output/predictions").outputs[0]

            # Generate batches for one epoch
            batches = pp.batch_iter(
                    list(x_test), FLAGS.batch_size, 1, shuffle=False)

            # Collect the predictions here
            all_predictions = []
            all_probabilities = None

            for x_test_batch in batches:
                batch_predictions_scores = sess.run(
                        [predictions, scores],
                        {input_x: x_test_batch, dropout_keep_prob: 1.0})
                all_predictions = np.concatenate(
                        [all_predictions, batch_predictions_scores[0]])
                probabilities = softmax(batch_predictions_scores[1])
                if all_probabilities is not None:
                    all_probabilities = np.concatenate(
                            [all_probabilities, probabilities])
                else:
                    all_probabilities = probabilities

    # Print accuracy if y_test is defined
    if y_test is not None:
        correct_predictions = float(sum(all_predictions == y_test))
        logger.debug("Total number of test examples: {}".format(len(y_test)))
        logger.info("")
        logger.info(
                "Accuracy: {:g}".format(
                        correct_predictions/float(len(y_test))))

        class_report = metrics.classification_report(
                y_test, all_predictions,
                target_names=datasets['target_names'])
        logger.info(class_report)

        confusion_matrix = ConfusionMatrix(y_test, all_predictions)
        logger.info(confusion_matrix)
        logger.info("")
        str_labels = "Labels : "
        for idx, label in enumerate(datasets['target_names']):
            str_labels += "{} = {}, ".format(idx, label)
        logger.info(str_labels)
        logger.info("")

    # Save the evaluation to a csv
    predictions_human_readable = np.column_stack(
            (np.array(x_raw),
             [int(prediction) for prediction in all_predictions],
             ["{}".format(probability) for probability in all_probabilities]))
    out_path = os.path.join(checkpoints_folder, "..", "prediction.csv")

    logger.info("Saving evaluation to {0}".format(out_path))

    with open(out_path, 'w') as f:
        csv.writer(f).writerows(predictions_human_readable)

    return (datasets['data'], all_predictions, datasets['target_names'],
            class_report)

if __name__ == '__main__':

    with open("config.yml", 'r') as ymlfile:
        cfg = yaml.load(ymlfile)

    # Parameters
    # ==================================================

    # Data Parameters

    # Eval Parameters
    tf.flags.DEFINE_integer("batch_size", 64, "Batch Size (default: 64)")
    tf.flags.DEFINE_string("checkpoint_dir", "",
                           "Checkpoint directory from training run")

    # Misc Parameters
    tf.flags.DEFINE_boolean("allow_soft_placement", True,
                            "Allow device soft device placement")
    tf.flags.DEFINE_boolean("log_device_placement", False,
                            "Log placement of ops on devices")
    tf.flags.DEFINE_boolean("aspects",
                            False,
                            "Scope widened to aspects and not only entities")

    # Precise if predictions is on features or polarity
    tf.flags.DEFINE_string("focus", "", "'feature' or 'polarity'")

    FLAGS = tf.flags.FLAGS
    FLAGS._parse_flags()

    # Logger
    # ==================================================

    logger = logging.getLogger()
    logger.setLevel(logging.DEBUG)

    # File handler which logs even debug messages
    file_handler = logging.FileHandler('log.log')
    file_handler.setLevel(logging.DEBUG)

    # Other file handler to store information for each run
    log_directory = os.path.join(FLAGS.checkpoint_dir, "eval.log")
    run_file_handler = logging.FileHandler(log_directory)
    run_file_handler.setLevel(logging.DEBUG)

    # Console handler which logs info messages
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)

    # Create formatter and add it to the handlers
    formatter = logging.Formatter("%(message)s")
    file_handler.setFormatter(formatter)
    run_file_handler.setFormatter(formatter)
    console_handler.setFormatter(formatter)

    # Add handlers to the logger
    logger.addHandler(file_handler)
    logger.addHandler(run_file_handler)
    logger.addHandler(console_handler)
    logger.debug(" *** Parameters *** ")
    for attr, value in sorted(FLAGS.__flags.items()):
        logger.debug("{}={}".format(attr.upper(), value))
    logger.debug("")

    # ----------
    # Prediction part :
    # ----------
    # First, construction of the representation of the XML document which we
    # want to predict --> Pandas.dataframe
    # Then, prediction of the outputs of both CNN_feature and CNN_polarity to
    # evaluate the accuracy of each CNN.
    # Afterwards, construction of the whole predictions.
    # Then, compute the accuracy, recall and f-score  to observ if the whole
    # model is good or not.
    # Finally, print the predicted and actual results in a CSV file.
    # ==================================================

    # ==================================================
    # Dataframe for actual and whole results
    # dataframe_actual = 'review_id', 'sentence_id', 'text', 'feature',
    #                    'polarity'
    # whole_prediction = 'review_id', 'sentence_id', 'text', 'feature',
    #                    'pred_feature', 'polarity', 'pred_polarity'
    # ==================================================

    dataset_name = cfg["datasets"]["default"]
    current_domain = cfg["datasets"][dataset_name]["current_domain"]
    if current_domain == 'RESTAURANT':
        dataframe_actual = pp.parse_XML(RESTAURANT_TEST, FLAGS.aspects)
        dataframe_actual = pp.select_and_simplify_dataset(
                dataframe_actual, RESTAURANT_TEST, FLAGS.aspects)
    elif current_domain == 'LAPTOP':
        dataframe_actual = pp.parse_XML(LAPTOP_TEST)
        dataframe_actual = pp.select_and_simplify_dataset(
                dataframe_actual, LAPTOP_TEST)
    else:
        raise ValueError("The 'current_domain' parameter in the " +
                         "'config.yml' file must be 'RESTAURANT' " +
                         "or 'LAPTOP'")

    whole_prediction = pd.DataFrame(data=None, columns=[
            'review_id', 'sentence_id', 'text', 'feature', 'pred_feature',
            'polarity', 'pred_polarity'])

    # ==================================================
    # CNN_feature predictions
    # ==================================================

    sentences_feature, all_predictions_feature, target_names_feature, feature_class_report =\
        prediction_process_CNN(FLAGS.checkpoint_dir, cfg, 'feature')

    # ==================================================
    # CNN_polarity predictions
    # ==================================================

    sentences_polarity, all_predictions_polarity, target_names_polarity, polarity_class_report =\
        prediction_process_CNN(FLAGS.checkpoint_dir, cfg, 'polarity')

    # ==================================================
    # Construction of the whole predictions
    # ==================================================
    for index, row in dataframe_actual.iterrows():
        review_id = row['review_id']
        sentence_id = row['sentence_id']
        text = row['text']
        feature = row['feature']
        polarity = row['polarity']

        # Feature
        # ==================================================

        # Retrieve index in the list of sentences
        index_text = sentences_feature.index(text)

        # Search the feature which corresponds to the text (retrieve the first
        # occurence)
        pred_feature = all_predictions_feature[index_text]

        # Translate to corresponding label
        pred_feature = target_names_feature[int(pred_feature)]

        # Polarity
        # ==================================================

        # Retrieve index in the list of sentences
        index_text = sentences_polarity.index(text)

        # Search the feature which corresponds to the text (retrieve the first
        # occurence)
        pred_polarity = all_predictions_polarity[index_text]

        # Translate to corresponding label
        pred_polarity = target_names_polarity[int(pred_polarity)]

        whole_prediction = whole_prediction.append(
                pd.DataFrame({'review_id': review_id,
                              'sentence_id': sentence_id,
                              'text': text,
                              'feature': feature,
                              'pred_feature': pred_feature,
                              'polarity': polarity,
                              'pred_polarity': pred_polarity},
                             index=[0]), ignore_index=True)

    # Add a column to check if the whole prediction is correct (feature and
    # pred_feature must be equal AND polarity and pred_polarity must also be
    # equal)
    whole_prediction['check'] =\
        ((whole_prediction.feature == whole_prediction.pred_feature) &
         (whole_prediction.polarity == whole_prediction.pred_polarity))

    # ==================================================
    # Effectiveness of the algorithm
    # ==================================================

    # Construction of dictionary to store new classes
    # Ex : FOOD, positive will be 0, FOOD, neutral : 1...etc...
    dict_polarity = {}
    for key, value in zip(POLARITY, list(range(len(POLARITY)))):
        dict_polarity[key] = value

    dict_entity_polarity = {}
    if current_domain == 'RESTAURANT':
        if not FLAGS.aspects:
            index = 0
            for entity in RESTAURANT_ENTITIES:
                dict_polarity = {}
                for polarity in POLARITY:
                    dict_polarity[polarity] = index
                    index += 1
                dict_entity_polarity[entity] = dict_polarity
        else:
            index = 0
            for entity in RESTAURANT_ASPECTS:
                dict_polarity = {}
                for polarity in POLARITY:
                    dict_polarity[polarity] = index
                    index += 1
                dict_entity_polarity[entity] = dict_polarity
    elif current_domain == 'LAPTOP':
        index = 0
        for entity in LAPTOP_ENTITIES:
            dict_polarity = {}
            for polarity in POLARITY:
                dict_polarity[polarity] = index
                index += 1
            dict_entity_polarity[entity] = dict_polarity
    else:
        raise ValueError("The 'current_domain' parameter in the " +
                         "'config.yml' file must be 'RESTAURANT' " +
                         "or 'LAPTOP'")

    # Create a new DataFrame to add to whole_prediction. The new DataFrame is
    # composed of 'new_class' and 'pred_new_class' columns
    list_of_rows = []
    for index, row in whole_prediction.iterrows():
        list_of_rows.append(
                [dict_entity_polarity[row['feature']][row['polarity']],
                 dict_entity_polarity[row['pred_feature']][row['pred_polarity']]])
    df_to_append = pd.DataFrame(data=list_of_rows,
                                columns=['new_class', 'pred_new_class'])
    whole_prediction = whole_prediction.assign(
             new_class=df_to_append['new_class'])
    whole_prediction = whole_prediction.assign(
             pred_new_class=df_to_append['pred_new_class'])

    logger.info("Effectiveness of the whole algorithm")
    logger.info("")
    class_report = metrics.classification_report(
            whole_prediction['new_class'],
            whole_prediction['pred_new_class'])
    logger.info(class_report)

    logger.info("")
    for entity, dict_polarity in dict_entity_polarity.items():
        for polarity, num_class in dict_polarity.items():
            logger.info("{} : {} - {}".format(num_class, entity, polarity))

    # Save the predictions into a CSV file inside the folder of the current run
    path_prediction_file = os.path.join(FLAGS.checkpoint_dir,
                                        'predictions.csv')
    whole_prediction.to_csv(path_prediction_file, encoding='utf-8',
                            columns=['review_id', 'sentence_id', 'text',
                                     'feature', 'pred_feature',
                                     'polarity', 'pred_polarity',
                                     'check', 'new_class', 'pred_new_class'])

    # ==================================================
    # Display charts
    # ==================================================
    an.bar_chart_classification_report(feature_class_report,
                                       "Effectiveness of CNN_feature",
                                       FLAGS.checkpoint_dir)
    an.bar_chart_classification_report(polarity_class_report,
                                       "Effectiveness of CNN_polarity",
                                       FLAGS.checkpoint_dir)
    an.bar_chart_classification_report(class_report,
                                       "Effectiveness of whole algorithm",
                                       FLAGS.checkpoint_dir)

    an.pie_chart_support_distribution(feature_class_report,
                                      "Data distribution for CNN_feature",
                                      FLAGS.checkpoint_dir)
    an.pie_chart_support_distribution(polarity_class_report,
                                      "Data distribution for CNN_polarity",
                                      FLAGS.checkpoint_dir)
    an.pie_chart_support_distribution(class_report,
                                      "Data distribution for whole algorithm",
                                      FLAGS.checkpoint_dir)


# This code is part of Mthree.
#
# (C) Copyright IBM 2021.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.
# pylint: disable=no-name-in-module, invalid-name
"""Calibration data"""

import warnings
import threading
import datetime
from time import perf_counter

import psutil
import numpy as np
import scipy.linalg as la
import scipy.sparse.linalg as spla
import orjson
from qiskit import QuantumCircuit, transpile, execute
from qiskit.providers import BaseBackend

from mthree.matrix import _reduced_cal_matrix, sdd_check
from mthree.utils import counts_to_vector, vector_to_quasiprobs
from mthree.norms import ainv_onenorm_est_lu, ainv_onenorm_est_iter
from mthree.matvec import M3MatVec
from mthree.exceptions import M3Error
from mthree.classes import QuasiCollection

def _tensor_meas_states(qubit, num_qubits, initial_reset=False):
    """Construct |0> and |1> states
    for independent 1Q cals.
    """
    qc0 = QuantumCircuit(num_qubits, 1)
    if initial_reset:
        qc0.reset(qubit)
    qc0.measure(qubit, 0)
    qc1 = QuantumCircuit(num_qubits, 1)
    if initial_reset:
        qc1.reset(qubit)
    qc1.x(qubit)
    qc1.measure(qubit, 0)
    return [qc0, qc1]

def _marg_meas_states(num_qubits, initial_reset=False):
    """Construct all zeros and all ones states
    for marginal 1Q cals.
    """
    qc0 = QuantumCircuit(num_qubits)
    if initial_reset:
        qc0.reset(range(num_qubits))
    qc0.measure_all()
    qc1 = QuantumCircuit(num_qubits)
    if initial_reset:
        qc1.reset(range(num_qubits))
    qc1.x(range(num_qubits))
    qc1.measure_all()
    return [qc0, qc1]

def _balanced_cal_strings(num_qubits):
    """Compute the 2*num_qubits strings for balanced calibration.

    Parameters:
        num_qubits (int): Number of qubits to be measured.

    Returns:
        list: List of strings for balanced calibration circuits.
    """
    strings = []
    for rep in range(1, num_qubits+1):
        str1 = ''
        str2 = ''
        for jj in range(int(np.ceil(num_qubits / rep))):
            str1 += str(jj % 2) * rep
            str2 += str((jj+1) % 2) * rep

        strings.append(str1[:num_qubits])
        strings.append(str2[:num_qubits])
    return strings

def _balanced_cal_circuits(cal_strings, initial_reset=False):
    """Build balanced calibration circuits.

    Parameters:
        cal_strings (list): List of strings for balanced cal circuits.

    Returns:
        list: List of balanced cal circuits.
    """
    num_qubits = len(cal_strings[0])
    circs = []
    for string in cal_strings:
        qc = QuantumCircuit(num_qubits)
        if initial_reset:
            qc.reset(range(num_qubits))
        for idx, bit in enumerate(string[::-1]):
            if bit == '1':
                qc.x(idx)
        qc.measure_all()
        circs.append(qc)
    return circs

class M3Mitigation():
    """Main M3 calibration class."""
    def __init__(self, system=None, iter_threshold=4096):
        """Main M3 calibration class.

        Parameters:
            system (BaseBackend): Target backend.
            iter_threshold (int): Sets the bitstring count at which iterative mode
                                  is turned on (assuming reasonable error rates).

        Attributes:
            system (BaseBackend): The target system.
            single_qubit_cals (list): 1Q calibration matrices
        """
        self.system = system
        self.single_qubit_cals = None
        self.num_qubits = system.configuration().num_qubits if system else None
        self.iter_threshold = iter_threshold
        self.cal_shots = None
        self.cal_method = 'balanced'
        self.cal_timestamp = None  # The time at which the cals result was generated
        self.rep_delay = None
        # attributes for handling threaded job
        self._thread = None
        self._job_error = None

    def __getattribute__(self, attr):
        """This allows for checking the status of the threaded cals call

        For certain attr this will join the thread and/or raise an error.
        """
        __dict__ = super().__getattribute__('__dict__')
        if attr in __dict__:
            if attr in ['single_qubit_cals']:
                self._thread_check()
        return super().__getattribute__(attr)

    def _form_cals(self, qubits):
        """Form the 1D cals array from tensored cals data

        Parameters:
            qubits (array_like): The qubits to calibrate over.

        Returns:
            ndarray: 1D Array of float cals data.
        """
        qubits = np.asarray(qubits, dtype=int)
        cals = np.zeros(4*qubits.shape[0], dtype=float)

        # Reverse index qubits for easier indexing later
        for kk, qubit in enumerate(qubits[::-1]):
            cals[4*kk:4*kk+4] = self.single_qubit_cals[qubit].ravel()
        return cals

    def _check_sdd(self, counts, qubits, distance=None):
        """Checks if reduced A-matrix is SDD or not

        Parameters:
            counts (dict): Dictionary of counts.
            qubits (array_like): List of qubits.
            distance (int): Distance to compute over.

        Returns:
            bool: True if A-matrix is SDD, else False

        Raises:
            M3Error: Number of qubits supplied does not match bit-string length.
        """
        # If distance is None, then assume max distance.
        num_bits = len(qubits)
        if distance is None:
            distance = num_bits

        # check if len of bitstrings does not equal number of qubits passed.
        bitstring_len = len(next(iter(counts)))
        if bitstring_len != num_bits:
            raise M3Error('Bitstring length ({}) does not match'.format(bitstring_len) +
                          ' number of qubits ({})'.format(num_bits))
        cals = self._form_cals(qubits)
        return sdd_check(counts, cals, num_bits, distance)

    def tensored_cals_from_system(self, qubits=None, shots=None,  method='balanced',
                                  rep_delay=None, cals_file=None):
        """Grab calibration data from system.

        Parameters:
            qubits (array_like): Qubits over which to correct calibration data. Default is all.
            shots (int): Number of shots per circuit. Default is min(1e4, max_shots).
            method (str): Type of calibration, 'balanced' (default), 'independent', or 'marginal'.
            rep_delay (float): Delay between circuits on IBM Quantum backends.
            cals_file (str): Output path to write JSON calibration data to.
        """
        warnings.warn("This method is deprecated, use 'cals_from_system' instead.")
        self.cals_from_system(qubits=qubits, shots=shots, method=method,
                              rep_delay=rep_delay,
                              cals_file=cals_file)

    def cals_from_system(self, qubits=None, shots=None, method='balanced',
                         initial_reset=False, rep_delay=None, cals_file=None):
        """Grab calibration data from system.

        Parameters:
            qubits (array_like): Qubits over which to correct calibration data. Default is all.
            shots (int): Number of shots per circuit. min(1e4, max_shots).
            method (str): Type of calibration, 'balanced' (default), 'independent', or 'marginal'.
            initial_reset (bool): Use resets at beginning of calibration circuits, default=False.
            rep_delay (float): Delay between circuits on IBM Quantum backends.
            cals_file (str): Output path to write JSON calibration data to.

        Raises:
            M3Error: Called while a calibration currently in progress.
        """
        if self._thread:
            raise M3Error('Calibration currently in progress.')
        if qubits is None:
            qubits = range(self.num_qubits)
        self.cal_method = method
        self.rep_delay = rep_delay
        self.cal_timestamp = None
        self._grab_additional_cals(qubits, shots=shots,  method=method,
                                   rep_delay=rep_delay, initial_reset=initial_reset)
        if cals_file:
            self.cals_to_file(cals_file)

    def cals_from_file(self, cals_file):
        """Generated the calibration data from a previous runs output

            cals_file (str): A string path to the saved counts file from an
                             earlier run.
            Raises:
                M3Error: Calibration in progress.
        """
        if self._thread:
            raise M3Error('Calibration currently in progress.')
        with open(cals_file, 'r', encoding='utf-8') as fd:
            loaded_data = orjson.loads(fd.read())
            if isinstance(loaded_data, dict):
                self.single_qubit_cals = [np.asarray(cal) if cal else None
                                          for cal in loaded_data['cals']]
                self.cal_timestamp = loaded_data['timestamp']
            else:
                warnings.warn('Loading from old M3 file format.  Save again to update.')
                self.cal_timestamp = None
                self.single_qubit_cals = [np.asarray(cal) if cal else None
                                          for cal in loaded_data]

    def cals_to_file(self, cals_file=None):
        """Save calibration data to JSON file.

            Parameters:
                cals_file (str): File in which to store calibrations.

            Raises:
                M3Error: Calibration filename missing.
                M3Error: Mitigator is not calibrated.
        """
        if not cals_file:
            raise M3Error('cals_file must be explicitly set.')
        if not self.single_qubit_cals:
            raise M3Error('Mitigator is not calibrated.')
        save_dict = {'timestamp': self.cal_timestamp,
                     'backend': self.system.name(),
                     'cals': self.single_qubit_cals}
        with open(cals_file, 'wb') as fd:
            fd.write(orjson.dumps(save_dict,
                                  option=orjson.OPT_SERIALIZE_NUMPY))

    def tensored_cals_from_file(self, cals_file):
        """Generated the tensored calibration data from a previous runs output

            cals_file (str): A string path to the saved counts file from an
                             earlier run.
        """
        warnings.warn("This method is deprecated, use 'cals_from_file' instead.")
        self.cals_from_file(cals_file)

    def _grab_additional_cals(self, qubits, shots=None, method='balanced', rep_delay=None,
                              initial_reset=False):
        """Grab missing calibration data from backend.

        Parameters:
            qubits (array_like): List of measured qubits.
            shots (int): Number of shots to take, min(1e4, max_shots).
            method (str): Type of calibration, 'balanced' (default), 'independent', or 'marginal'.
            rep_delay (float): Delay between circuits on IBM Quantum backends.
            initial_reset (bool): Use resets at beginning of calibration circuits, default=False.

        Raises:
            M3Error: Backend not set.
            M3Error: Faulty qubits found.
        """
        if self.system is None:
            raise M3Error("System is not set.  Use 'cals_from_file'.")
        if self.single_qubit_cals is None:
            self.single_qubit_cals = [None]*self.num_qubits
        if self.cal_shots is None:
            if shots is None:
                shots = min(self.system.configuration().max_shots, 10000)
            self.cal_shots = shots
        if self.rep_delay is None:
            self.rep_delay = rep_delay

        if method not in ['independent', 'balanced', 'marginal']:
            raise M3Error('Invalid calibration method.')

        if isinstance(qubits, dict):
            # Assuming passed a mapping
            qubits = list(qubits)
        elif isinstance(qubits, list):
            # Check if passed a list of mappings
            if isinstance(qubits[0], dict):
                # Assuming list of mappings, need to get unique elements
                _qubits = []
                for item in qubits:
                    _qubits.extend(list(item))
                qubits = list(set(_qubits))

        num_cal_qubits = len(qubits)
        cal_strings = []
        if method == 'marginal':
            circs = _marg_meas_states(num_cal_qubits, initial_reset=initial_reset)
            trans_qcs = transpile(circs, self.system,
                                  initial_layout=qubits, optimization_level=0)
        elif method == 'balanced':
            cal_strings = _balanced_cal_strings(num_cal_qubits)
            circs = _balanced_cal_circuits(cal_strings, initial_reset=initial_reset)
            trans_qcs = transpile(circs, self.system,
                                  initial_layout=qubits, optimization_level=0)
        # Indeopendent
        else:
            circs = []
            for kk in qubits:
                circs.extend(_tensor_meas_states(kk, self.num_qubits,
                                                 initial_reset=initial_reset))
            trans_qcs = transpile(circs, self.system, optimization_level=0)

        # This BaseBackend check is here for Qiskit direct access.  Should be removed later.
        if isinstance(self.system, BaseBackend):
            job = execute(trans_qcs, self.system, optimization_level=0,
                          shots=self.cal_shots, rep_delay=self.rep_delay)
        else:
            job = self.system.run(trans_qcs, shots=self.cal_shots, rep_delay=self.rep_delay)

        # Execute job and cal building in new theread.
        self._job_error = None
        thread = threading.Thread(target=_job_thread, args=(job, self, method, qubits,
                                                            num_cal_qubits, cal_strings))
        self._thread = thread
        self._thread.start()

    def apply_correction(self, counts, qubits, distance=None,
                         method='auto',
                         max_iter=25, tol=1e-5,
                         return_mitigation_overhead=False,
                         details=False):
        """Applies correction to given counts.

        Parameters:
            counts (dict, list): Input counts dict or list of dicts.
            qubits (array_like): Qubits on which measurements applied.
            distance (int): Distance to correct for. Default=num_bits
            method (str): Solution method: 'auto', 'direct' or 'iterative'.
            max_iter (int): Max. number of iterations, Default=25.
            tol (float): Convergence tolerance of iterative method, Default=1e-5.
            return_mitigation_overhead (bool): Returns the mitigation overhead, default=False.
            details (bool): Return extra info, default=False.

        Returns:
            QuasiDistribution or QuasiCollection: Dictionary of quasiprobabilities if
                                                  input is a single dict, else a collection
                                                  of quasiprobabilities.

        Raises:
            M3Error: Bitstring length does not match number of qubits given.
        """
        if len(counts) == 0:
            raise M3Error('Input counts is any empty dict.')
        given_list = False
        if isinstance(counts, (list, np.ndarray)):
            given_list = True
        if not given_list:
            counts = [counts]

        if isinstance(qubits, dict):
            # If a mapping was given for qubits
            qubits = [list(qubits)]
        elif not any(isinstance(qq, (list, tuple, np.ndarray, dict)) for qq in qubits):
            qubits = [qubits]*len(counts)
        else:
            if isinstance(qubits[0], dict):
                # assuming passed a list of mappings
                qubits = [list(qu) for qu in qubits]

        if len(qubits) != len(counts):
            raise M3Error('Length of counts does not match length of qubits.')

        quasi_out = []
        for idx, cnts in enumerate(counts):

            quasi_out.append(
                self._apply_correction(cnts, qubits=qubits[idx],
                                       distance=distance,
                                       method=method,
                                       max_iter=max_iter, tol=tol,
                                       return_mitigation_overhead=return_mitigation_overhead,
                                       details=details)
                            )

        if not given_list:
            return quasi_out[0]
        return QuasiCollection(quasi_out)

    def _apply_correction(self, counts, qubits, distance=None,
                          method='auto',
                          max_iter=25, tol=1e-5,
                          return_mitigation_overhead=False,
                          details=False):
        """Applies correction to given counts.

        Parameters:
            counts (dict): Input counts dict.
            qubits (array_like): Qubits on which measurements applied.
            distance (int): Distance to correct for. Default=num_bits
            method (str): Solution method: 'auto', 'direct' or 'iterative'.
            max_iter (int): Max. number of iterations, Default=25.
            tol (float): Convergence tolerance of iterative method, Default=1e-5.
            return_mitigation_overhead (bool): Returns the mitigation overhead, default=False.
            details (bool): Return extra info, default=False.

        Returns:
            QuasiDistribution: Dictionary of quasiprobabilities.

        Raises:
            M3Error: Bitstring length does not match number of qubits given.
        """
        # This is needed because counts is a Counts object in Qiskit not a dict.
        counts = dict(counts)
        shots = sum(counts.values())

        # If distance is None, then assume max distance.
        num_bits = len(qubits)
        num_elems = len(counts)
        if distance is None:
            distance = num_bits

        # check if len of bitstrings does not equal number of qubits passed.
        bitstring_len = len(next(iter(counts)))
        if bitstring_len != num_bits:
            raise M3Error('Bitstring length ({}) does not match'.format(bitstring_len) +
                          ' number of qubits ({})'.format(num_bits))

        # Check if no cals done yet
        if self.single_qubit_cals is None:
            warnings.warn('No calibration data. Calibrating: {}'.format(qubits))
            self._grab_additional_cals(qubits, method=self.cal_method)

        # Check if one or more new qubits need to be calibrated.
        missing_qubits = [qq for qq in qubits if self.single_qubit_cals[qq] is None]
        if any(missing_qubits):
            warnings.warn('Computing missing calibrations for qubits: {}'.format(missing_qubits))
            self._grab_additional_cals(missing_qubits, method=self.cal_method)

        if method == 'auto':
            current_free_mem = psutil.virtual_memory().available / 1024**3
            # First check if direct method can be run
            if num_elems <= self.iter_threshold \
                    and ((num_elems**2+num_elems)*8/1024**3 < current_free_mem/2):
                method = 'direct'
            else:
                method = 'iterative'

        if method == 'direct':
            st = perf_counter()
            mit_counts, col_norms, gamma = self._direct_solver(counts, qubits, distance,
                                                               return_mitigation_overhead)
            dur = perf_counter()-st
            mit_counts.shots = shots
            if gamma is not None:
                mit_counts.mitigation_overhead = gamma * gamma
            if details:
                info = {'method': 'direct', 'time': dur, 'dimension': num_elems}
                info['col_norms'] = col_norms
                return mit_counts, info
            return mit_counts

        elif method == 'iterative':
            iter_count = np.zeros(1, dtype=int)

            def callback(_):
                iter_count[0] += 1

            if details:
                st = perf_counter()
                mit_counts, col_norms, gamma = self._matvec_solver(counts,
                                                                   qubits,
                                                                   distance,
                                                                   tol,
                                                                   max_iter,
                                                                   1,
                                                                   callback,
                                                                   return_mitigation_overhead)
                dur = perf_counter()-st
                mit_counts.shots = shots
                if gamma is not None:
                    mit_counts.mitigation_overhead = gamma * gamma
                info = {'method': 'iterative', 'time': dur, 'dimension': num_elems}
                info['iterations'] = iter_count[0]
                info['col_norms'] = col_norms
                return mit_counts, info
            # pylint: disable=unbalanced-tuple-unpacking
            mit_counts, gamma = self._matvec_solver(counts, qubits, distance, tol,
                                                    max_iter, 0, None,
                                                    return_mitigation_overhead)
            mit_counts.shots = shots
            if gamma is not None:
                mit_counts.mitigation_overhead = gamma * gamma
            return mit_counts

        else:
            raise M3Error('Invalid method: {}'.format(method))

    def reduced_cal_matrix(self, counts, qubits, distance=None):
        """Return the reduced calibration matrix used in the solution.

        Parameters:
            counts (dict): Input counts dict.
            qubits (array_like): Qubits on which measurements applied.
            distance (int): Distance to correct for. Default=num_bits

        Returns:
            ndarray: 2D array of reduced calibrations.
            dict: Counts in order they are displayed in matrix.

        Raises:
            M3Error: If bit-string length does not match passed number
                     of qubits.
        """
        counts = dict(counts)
        # If distance is None, then assume max distance.
        num_bits = len(qubits)
        if distance is None:
            distance = num_bits

        # check if len of bitstrings does not equal number of qubits passed.
        bitstring_len = len(next(iter(counts)))
        if bitstring_len != num_bits:
            raise M3Error('Bitstring length ({}) does not match'.format(bitstring_len) +
                          ' number of qubits ({})'.format(num_bits))

        cals = self._form_cals(qubits)
        A, counts, _ = _reduced_cal_matrix(counts, cals, num_bits, distance)
        return A, counts

    def _direct_solver(self, counts, qubits, distance=None,
                       return_mitigation_overhead=False):
        """Apply the mitigation using direct LU factorization.

        Parameters:
            counts (dict): Input counts dict.
            qubits (int): Qubits over which to calibrate.
            distance (int): Distance to correct for. Default=num_bits
            return_mitigation_overhead (bool): Returns the mitigation overhead, default=False.

        Returns:
            QuasiDistribution: dict of Quasiprobabilites
        """
        cals = self._form_cals(qubits)
        num_bits = len(qubits)
        A, sorted_counts, col_norms = _reduced_cal_matrix(counts, cals, num_bits, distance)
        vec = counts_to_vector(sorted_counts)
        LU = la.lu_factor(A, check_finite=False)
        x = la.lu_solve(LU, vec, check_finite=False)
        gamma = None
        if return_mitigation_overhead:
            gamma = ainv_onenorm_est_lu(A, LU)
        out = vector_to_quasiprobs(x, sorted_counts)
        return out, col_norms, gamma

    def _matvec_solver(self, counts, qubits, distance, tol=1e-5, max_iter=25,
                       details=0, callback=None,
                       return_mitigation_overhead=False):
        """Compute solution using GMRES and Jacobi preconditioning.

        Parameters:
            counts (dict): Input counts dict.
            qubits (int): Qubits over which to calibrate.
            tol (float): Tolerance to use.
            max_iter (int): Maximum number of iterations to perform.
            distance (int): Distance to correct for. Default=num_bits
            details (bool): Return col norms.
            callback (callable): Callback function to record iteration count.
            return_mitigation_overhead (bool): Returns the mitigation overhead, default=False.

        Returns:
            QuasiDistribution: dict of Quasiprobabilites

        Raises:
            M3Error: Solver did not converge.
        """
        cals = self._form_cals(qubits)
        M = M3MatVec(dict(counts), cals, distance)
        L = spla.LinearOperator((M.num_elems, M.num_elems),
                                matvec=M.matvec, rmatvec=M.rmatvec)
        diags = M.get_diagonal()

        def precond_matvec(x):
            out = x / diags
            return out

        P = spla.LinearOperator((M.num_elems, M.num_elems), precond_matvec)
        vec = counts_to_vector(M.sorted_counts)
        out, error = spla.gmres(L, vec, tol=tol, atol=tol, maxiter=max_iter,
                                M=P, callback=callback)
        if error:
            raise M3Error('GMRES did not converge: {}'.format(error))

        gamma = None
        if return_mitigation_overhead:
            gamma = ainv_onenorm_est_iter(M, tol=tol, max_iter=max_iter)

        quasi = vector_to_quasiprobs(out, M.sorted_counts)
        if details:
            return quasi, M.get_col_norms(), gamma
        return quasi, gamma

    def readout_fidelity(self, qubits=None):
        """Compute readout fidelity for calibrated qubits.

        Parameters:
            qubits (array_like): Qubits to compute over, default is all.

        Returns:
            list: List of qubit fidelities.

        Raises:
            M3Error: Mitigator is not calibrated.
            M3Error: Qubit indices out of range.
        """
        if self.single_qubit_cals is None:
            raise M3Error('Mitigator is not calibrated')

        if qubits is None:
            qubits = range(self.num_qubits)
        else:
            outliers = [kk for kk in qubits if kk >= self.num_qubits]
            if any(outliers):
                raise M3Error('One or more qubit indices out of range: {}'.format(outliers))
        fids = []
        for kk in qubits:
            qubit = self.single_qubit_cals[kk]
            if qubit is not None:
                fids.append(np.mean(qubit.diagonal()))
            else:
                fids.append(None)
        return fids

    def _thread_check(self):
        """Check if a thread is running and join it.

        Raise an error if one is given.
        """
        if self._thread and self._thread != threading.current_thread():
            self._thread.join()
            self._thread = None
        if self._job_error:
            raise self._job_error  # pylint: disable=raising-bad-type

def _job_thread(job, mit, method, qubits, num_cal_qubits, cal_strings):
    """Run the calibration job in a different thread and post-process

    Parameters:
        mit (M3Mitigator): The mitigator instance
        method (str): The type of calibration
        qubits (list): List of qubits used
        num_cal_qubits (int): Number of calibration qubits
        cal_strings (list): List of cal strings for balanced cals
    """
    try:
        res = job.result()
    # pylint: disable=broad-except
    except Exception as error:
        mit._job_error = error
        return
    counts = res.get_counts()
    # attach timestamp
    timestamp = res.date
    # Needed since Aer result date is str but IBMQ job is datetime
    if isinstance(timestamp, datetime.datetime):
        timestamp = timestamp.isoformat()
    mit.cal_timestamp = timestamp
    # A list of qubits with bad meas cals
    bad_list = []
    if method == 'independent':
        for idx, qubit in enumerate(qubits):
            mit.single_qubit_cals[qubit] = np.zeros((2, 2), dtype=float)
            # Counts 0 has all P00, P10 data, so do that here
            prep0_counts = counts[2*idx]
            P10 = prep0_counts.get('1', 0) / mit.cal_shots
            P00 = 1-P10
            mit.single_qubit_cals[qubit][:, 0] = [P00, P10]
            # plus 1 here since zeros data at pos=0
            prep1_counts = counts[2*idx+1]
            P01 = prep1_counts.get('0', 0) / mit.cal_shots
            P11 = 1-P01
            mit.single_qubit_cals[qubit][:, 1] = [P01, P11]
            if P01 >= P00:
                bad_list.append(qubit)
    elif method == 'marginal':
        prep0_counts = counts[0]
        prep1_counts = counts[1]
        for idx, qubit in enumerate(qubits):
            mit.single_qubit_cals[qubit] = np.zeros((2, 2), dtype=float)
            count_vals = 0
            index = num_cal_qubits-idx-1
            for key, val in prep0_counts.items():
                if key[index] == '0':
                    count_vals += val
            P00 = count_vals / mit.cal_shots
            P10 = 1-P00
            mit.single_qubit_cals[qubit][:, 0] = [P00, P10]
            count_vals = 0
            for key, val in prep1_counts.items():
                if key[index] == '1':
                    count_vals += val
            P11 = count_vals / mit.cal_shots
            P01 = 1-P11
            mit.single_qubit_cals[qubit][:, 1] = [P01, P11]
            if P01 >= P00:
                bad_list.append(qubit)
    # balanced calibration
    else:
        cals = [np.zeros((2, 2), dtype=float) for kk in range(num_cal_qubits)]

        for idx, count in enumerate(counts):

            target = cal_strings[idx][::-1]
            good_prep = np.zeros(num_cal_qubits, dtype=float)
            denom = mit.cal_shots * num_cal_qubits

            for key, val in count.items():
                key = key[::-1]
                for kk in range(num_cal_qubits):
                    if key[kk] == target[kk]:
                        good_prep[kk] += val

            for kk, cal in enumerate(cals):
                if target[kk] == '0':
                    cal[0, 0] += good_prep[kk] / denom
                else:
                    cal[1, 1] += good_prep[kk] / denom

        for jj, cal in enumerate(cals):
            cal[1, 0] = 1.0 - cal[0, 0]
            cal[0, 1] = 1.0 - cal[1, 1]

            if cal[0, 1] >= cal[0, 0]:
                bad_list.append(qubits[jj])

        for idx, cal in enumerate(cals):
            mit.single_qubit_cals[qubits[idx]] = cal

    if any(bad_list):
        mit._job_error = M3Error('Faulty qubits detected: {}'.format(bad_list))


from __future__ import absolute_import, division, print_function

import os
import math
import numpy as np

import torch
import torch.nn as nn
import torch.legacy.optim as legacyOptim

import torch.nn.functional as F
import torch.multiprocessing as mp
from torchvision import datasets, transforms

from model import ES

#import matplotlib.pyplot as plt

def do_rollouts(args, models, random_seeds, return_queue, data, label, are_negative):
    """
    For each model, do a rollout. Supports multiple models per thread but
    don't do it -- it's inefficient (it's mostly a relic of when I would run
    both a perturbation and its antithesis on the same thread).
    """
    all_returns = []
    all_acc = []
    for model in models:
        # Rollout
        logit = model(data)
        with torch.no_grad():
            reward = -F.cross_entropy(logit, label).numpy()
            _, preds = torch.max(logit, 1)
            all_acc.append(torch.sum(preds == label.data).numpy() / len(label))
        all_returns.append(reward)
    return_queue.put((random_seeds, all_returns, all_acc, are_negative))

def perturb_model(args, model, random_seed):
    """
    Modifies the given model with a pertubation of its parameters,
    as well as the negative perturbation, and returns both perturbed
    models.
    """
    new_model = ES(args.small_net)
    anti_model = ES(args.small_net)
    new_model.load_state_dict(model.state_dict())
    anti_model.load_state_dict(model.state_dict())
    np.random.seed(random_seed)
    for (k, v), (anti_k, anti_v) in zip(new_model.es_params(),
                                        anti_model.es_params()):
        eps = np.random.normal(0, 1, v.size())
        v += torch.from_numpy(args.sigma*eps).float()
        anti_v += torch.from_numpy(args.sigma*-eps).float()
    return [new_model, anti_model]

optimConfig = []
averageReward = []
maxReward = []
minReward = []
episodeCounter = []

def gradient_update(args, synced_model, returns, accs, random_seeds, neg_list,
                    num_eps, chkpt_dir, unperturbed_results):
    def fitness_shaping(returns):
        """
        A rank transformation on the rewards, which reduces the chances
        of falling into local optima early in training.
        """
        sorted_returns_backwards = sorted(returns)[::-1]
        lamb = len(returns)
        shaped_returns = []
        denom = sum([max(0, math.log(lamb/2 + 1, 2) -
                         math.log(sorted_returns_backwards.index(r) + 1, 2))
                     for r in returns])
        for r in returns:
            num = max(0, math.log(lamb/2 + 1, 2) -
                      math.log(sorted_returns_backwards.index(r) + 1, 2))
            shaped_returns.append(num/denom + 1/lamb)
        return shaped_returns

    def unperturbed_rank(returns, unperturbed_results):
        nth_place = 1
        for r in returns:
            if r > unperturbed_results:
                nth_place += 1
        rank_diag = ('%d out of %d (1 means gradient '
                     'is uninformative)' % (nth_place,
                                             len(returns) + 1))
        return rank_diag, nth_place

    batch_size = len(returns)
    assert batch_size == args.n
    assert len(random_seeds) == batch_size
    shaped_returns = fitness_shaping(returns)
    rank_diag, rank = unperturbed_rank(returns, unperturbed_results)
    if not args.silent:
        print('Episode num: %d\n'
              'Average reward: %f\n'
              'Variance in rewards: %f\n'
              'Max accuracy: %f\n'
              'Max reward: %f\n'
              'Min reward: %f\n'
              'Batch size: %d\n'
              'Max episode length: %d\n'
              'Sigma: %f\n'
              'Learning rate: %f\n'
              'Unperturbed reward: %f\n'
              'Unperturbed rank: %s\n'
              'Using Adam: %r\n' %
              (num_eps, np.mean(returns), np.var(returns), max(accs), max(returns),
               min(returns), batch_size,
               args.max_episode_length, args.sigma, args.lr,
               unperturbed_results, rank_diag, args.useAdam))

    averageReward.append(np.mean(returns))
    episodeCounter.append(num_eps)
    maxReward.append(max(returns))
    minReward.append(min(returns))

#    pltAvg, = plt.plot(episodeCounter, averageReward, label='average')
#    pltMax, = plt.plot(episodeCounter, maxReward,  label='max')
#    pltMin, = plt.plot(episodeCounter, minReward,  label='min')

#    plt.ylabel('rewards')
#    plt.xlabel('episode num')
#    plt.legend(handles=[pltAvg, pltMax,pltMin])

#    fig1 = plt.gcf()

#    plt.draw()
#    fig1.savefig('graph.png', dpi=100)

    # For each model, generate the same random numbers as we did
    # before, and update parameters. We apply weight decay once.
    if args.useAdam:
        globalGrads = None
        for i in range(args.n):
            np.random.seed(random_seeds[i])
            multiplier = -1 if neg_list[i] else 1
            reward = shaped_returns[i]

            localGrads = []
            idx = 0
            for k, v in synced_model.es_params():
                eps = np.random.normal(0, 1, v.size())
                grad = torch.from_numpy((args.n*args.sigma) * (reward*multiplier*eps)).float()

                localGrads.append(grad)

                if len(optimConfig) == idx:
                    optimConfig.append({ 'learningRate' : args.lr  })
                idx = idx + 1

            if globalGrads == None:
                globalGrads = localGrads
            else:
                for i in range(len(globalGrads)):
                    globalGrads[i] = torch.add(globalGrads[i], localGrads[i])

        idx = 0
        for k, v in synced_model.es_params():
            r, _ = legacyOptim.adam( lambda x:  (1, -globalGrads[idx]), v , optimConfig[idx])
            v.copy_(r)
            idx = idx + 1
    else:
        # For each model, generate the same random numbers as we did
        # before, and update parameters. We apply weight decay once.
        for i in range(args.n):
            np.random.seed(random_seeds[i])
            multiplier = -1 if neg_list[i] else 1
            reward = shaped_returns[i]
            for k, v in synced_model.es_params():
                eps = np.random.normal(0, 1, v.size())
                v += torch.from_numpy(args.lr/(args.n*args.sigma) *
                                      (reward*multiplier*eps)).float()
        args.lr *= args.lr_decay

    torch.save(synced_model.state_dict(),
               os.path.join(chkpt_dir, 'latest.pth'))
    return synced_model

def render_env(args, model):
    return

def generate_seeds_and_models(args, synced_model):
    """
    Returns a seed and 2 perturbed models
    """
    np.random.seed()
    random_seed = np.random.randint(2**30)
    two_models = perturb_model(args, synced_model, random_seed)
    return random_seed, two_models

def train_loop(args, synced_model, chkpt_dir):
    def flatten(raw_results, index):
        notflat_results = [result[index] for result in raw_results]
        return [item for sublist in notflat_results for item in sublist]
    print("Num params in network %d" % synced_model.count_parameters())
    num_eps = 0
    total_num_frames = 0

    transform = transforms.Compose([transforms.ToTensor(),
                               transforms.Normalize((0.1307,), (0.3081,))])
    data_train = datasets.MNIST(root = "./data/",
                            transform=transform,
                            train = True)
    data_test = datasets.MNIST(root="./data/",
                           transform = transform,
                           train = False)
    dataloaders = {"train": torch.utils.data.DataLoader(dataset=data_train,
                                                batch_size = 64,
                                                shuffle = True),
                   "test": torch.utils.data.DataLoader(dataset=data_test,
                                               batch_size = 64,
                                               shuffle = True)}
    steps = 0
    for _ in range(args.max_gradient_updates):
        for data, label in dataloaders["train"]:
            processes = []
            return_queue = mp.Queue()
            all_seeds, all_models = [], []
            # Generate a perturbation and its antithesis
            for j in range(int(args.n/2)):
                random_seed, two_models = generate_seeds_and_models(args,
                                                                synced_model)
            # Add twice because we get two models with the same seed
                all_seeds.append(random_seed)
                all_seeds.append(random_seed)
                all_models += two_models
            assert len(all_seeds) == len(all_models)
        # Keep track of which perturbations were positive and negative
        # Start with negative true because pop() makes us go backwards
            is_negative = True
        # Add all peturbed models to the queue
            while all_models:
                perturbed_model = all_models.pop()
                seed = all_seeds.pop()
                p = mp.Process(target=do_rollouts, args=(args,
                                                     [perturbed_model],
                                                     [seed],
                                                     return_queue,
                                                     data,
                                                     label,
                                                     [is_negative]))
                p.start()
                processes.append(p)
                is_negative = not is_negative
            assert len(all_seeds) == 0
        # Evaluate the unperturbed model as well
            p = mp.Process(target=do_rollouts, args=(args, [synced_model],
                                                 ['dummy_seed'],
                                                 return_queue, data, label,
                                                 ['dummy_neg']))
            p.start()
            processes.append(p)
            for p in processes:
                p.join()
            raw_results = [return_queue.get() for p in processes]
            seeds, results, accs, neg_list = [flatten(raw_results, index)
                                                for index in [0, 1, 2, 3]]
        # Separate the unperturbed results from the perturbed results
            _ = unperturbed_index = seeds.index('dummy_seed')
            seeds.pop(unperturbed_index)
            unperturbed_results = results.pop(unperturbed_index)
            _ = neg_list.pop(unperturbed_index)

            num_eps += len(results)
            synced_model = gradient_update(args, synced_model, results, accs, seeds,
                                       neg_list, num_eps,
                                       chkpt_dir, unperturbed_results)
#            steps += 1
#            if steps % 10 == 0:
#                running_corrects = 0
#                synced_model.eval()
#                for inputs, labels in dataloaders["test"]:
#                    with torch.no_grad():
#                    outputs = synced_model(inputs)
#                    _, preds = torch.max(outputs, 1)
#                    running_corrects += torch.sum(preds == labels.data)
#                epoch_acc = running_corrects.double() / len(data_test)
#                print('Step {} Acc: {:.4f}'.format(steps, epoch_acc))
#                synced_model.train()


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# Copyright (C) 2020  The Project U-Ray Authors.
#
# Use of this source code is governed by a ISC-style
# license that can be found in the LICENSE file or at
# https://opensource.org/licenses/ISC
#
# SPDX-License-Identifier: ISC

import csv
import numpy as np
from utils import util
from utils.clock_utils import make_bufg, MAX_GLOBAL_CLOCKS
from prjuray.db import Database

CMT_XY_FUN = util.create_xy_fun('')

def print_top(seed):
    np.random.seed(seed)

    db = Database(util.get_db_root(), util.get_part())
    grid = db.grid()

    site_to_tile = {}
    site_to_site_type = {}
    for tile_name in sorted(grid.tiles()):
        loc = grid.loc_of_tilename(tile_name)
        gridinfo = grid.gridinfo_at_loc(loc)

        for site, site_type in gridinfo.sites.items():
            site_to_site_type[site] = site_type
            site_to_tile[site] = tile_name

    bufce_row_drivers = {}

    with open('../active_bufce_row.csv') as f:
        for row in csv.DictReader(f):
            tile = site_to_tile[row['site']]

            gridinfo = grid.gridinfo_at_tilename(tile)
            if gridinfo.tile_type == 'CMT_RIGHT':
                continue

            clock_region = CMT_XY_FUN(row['site_clock_region'])

            key = clock_region, int(row['hdistr_number'])
            assert key not in bufce_row_drivers
            bufce_row_drivers[key] = row['site']

    bufg_drivers = {}

    with open('../bufg_outputs.csv') as f:
        for row in csv.DictReader(f):
            if site_to_site_type[row['site']] != 'BUFGCE_HDIO':
                continue

            clock_region = CMT_XY_FUN(row['clock_region'])
            if clock_region not in bufg_drivers:
                bufg_drivers[clock_region] = []

            bufg_drivers[clock_region].append(row['site'])

    for clock_region in bufg_drivers:
        bufg_drivers[clock_region].sort()
        np.random.shuffle(bufg_drivers[clock_region])

    clock_regions = set()
    for clock_region, _ in bufce_row_drivers.keys():
        clock_regions.add(clock_region)

    for clock_region in bufg_drivers.keys():
        clock_regions.add(clock_region)

    print("""
module top();
        """)

    bufg_routes = []
    for hroute in range(MAX_GLOBAL_CLOCKS):
        if np.random.randint(2):
            continue

        select_clock_regions = sorted(clock_regions)
        np.random.shuffle(select_clock_regions)

        bufg = None
        while len(select_clock_regions) > 0:
            clock_region = select_clock_regions.pop()
            if clock_region not in bufg_drivers:
                continue

            if len(bufg_drivers[clock_region]) > 0:
                bufg = bufg_drivers[clock_region].pop()
                break

        if bufg is None:
            continue

        sel_x, sel_y = clock_region
        clock_regions_for_bufg = []
        for x, y in clock_regions:
            if y == sel_y:
                clock_regions_for_bufg.append((x, y))

        clock_regions_for_bufg.sort()
        np.random.shuffle(clock_regions_for_bufg)

        bufce_row = None
        while True:
            clock_region = clock_regions_for_bufg.pop()
            key = clock_region, hroute

            if key in bufce_row_drivers:
                bufce_row = bufce_row_drivers[key]
                break

        if bufce_row is None:
            continue

        bufg_s, bufg_o_wire = make_bufg(
            site=bufg,
            site_type=site_to_site_type[bufg],
            idx=hroute,
            ce_inputs=['1'],
            randlib=np.random)

        print(bufg_s)

        print("""
        (* LOC="{loc}", KEEP, DONT_TOUCH *) BUFCE_ROW row{idx} (
            .I({bufg_o_wire})
        );""".format(
            loc=bufce_row,
            idx=hroute,
            bufg_o_wire=bufg_o_wire,
        ))

        bufg_routes.append((bufg, bufce_row, hroute))

    print('endmodule')

    with open('complete_top.tcl', 'w') as f:
        print(
            """
source $::env(URAY_UTILS_DIR)/utils.tcl
write_checkpoint -force design_reset.dcp
close_design
open_checkpoint design_reset.dcp
""",
            file=f)

        for bufg, bufce_row, hroute in bufg_routes:
            if np.random.randint(2):
                print(
                    "route_bufg_to_bufce_row_via_opposite_wire {bufg} {bufce_row} {hroute}"
                    .format(bufg=bufg, bufce_row=bufce_row, hroute=hroute),
                    file=f)

        print("""
route_design
""", file=f)