#!/usr/bin/env python3

#
# Copyright 2015 Dovetail Genomics LLC
#
#

from __future__ import division
from __future__ import print_function
from builtins import map
from builtins import str
from builtins import range
from past.utils import old_div
import sys
import networkx as nx
import chicago_edge_scores as ces

default_gapsize=100.0

ll={}
links={}

def is_shaved_tail(G,shave_round,shaved_degree,shave_limit):
    leaves=[]
    r={}
    for n in G.nodes():
        if shave_round.get(n,0)>shave_limit and shaved_degree.get(n,0)==1:
            leaves.append(n)
    for l in leaves:
        q=[l]
        while len(q)>0:
            n=q.pop()
            r[n]=True
            for nn in G.neighbors(n):
                if shave_round.get(nn,0)>shave_limit and shaved_degree.get(nn,0)<=2 and (not nn in q) and (not nn in r ):
                    q.append(nn)
    return r

def distance_to_nearest_branch(G,shave_round,shave_limit,trim_degree):
    branchpoints=[]
    r={}

    for n in G.nodes():
        #print "#",n, G.degree(n),G.degree(n)==1
        if shave_round.get(n,0)>shave_limit and trim_degree.get(n,0)>2:
            branchpoints.append(n)
#    print leaves

    round=1
#    print "#",leaves
    boundary=list(branchpoints)
    next=[]
    done=[]
    #r={}
    for b in boundary: r[b]=round
    while len(boundary)>0:
        round +=1
        next=[]
        for n in boundary:
            for nn in G.neighbors(n):
                if (not nn in boundary+done+next) :
                    next.append(nn)
#                    r[nn]=round
        for b in next:
            r[b] = round
#            if len( [nn for nn in G.neighbors(b) if not r.has_key(nn)] )==1: r[b]=round
        for b in boundary: done.append(b)
        boundary = []
        for n in next:
            if n in r: boundary.append(n)
#        next=[]

#    for n in G.nodes():
#        if not r.has_key(n): r[n]=round
    return r

def distance_to_nearest_leaf(G,shave_round,shave_limit,trim_degree):
    leaves=[]
    r={}
    for n in G.nodes():
        #print "#",n, G.degree(n),G.degree(n)==1
        if shave_round.get(n,0)<=shave_limit :
            r[n]=0
    for n in G.nodes():
        #print "#",n, G.degree(n),G.degree(n)==1
        if trim_degree.get(n,0)==1 and shave_round.get(n,0)>shave_limit:
            leaves.append(n)
#    print leaves

    round=1
#    print "#",leaves
    boundary=list(leaves)
    next=[]
    done=[]
#    r={}
    for b in boundary: r[b]=round
    while len(boundary)>0:
        round +=1
        next=[]
        for n in boundary:
            for nn in G.neighbors(n):
                if (not nn in boundary+done+next):
                    next.append(nn)
#                    r[nn]=round
        for b in next:
            r[b] = round
#            if len( [nn for nn in G.neighbors(b) if not r.has_key(nn)] )==1: r[b]=round
        for b in boundary: done.append(b)
        boundary = []
        for n in next:
            if n in r: boundary.append(n)
#        next=[]
#    for n in G.nodes():
#        if not r.has_key(n): r[n]=round
    return r

def shave_round(G):
    leaves=[]
    for n in G.nodes():
        #print "#",n, G.degree(n),G.degree(n)==1
        if G.degree(n)==1:
            leaves.append(n)
#    print leaves

    round=1
#    print "#",leaves
    boundary=list(leaves)
    next=[]
    done=[]
    r={}
    for b in boundary: r[b]=round
    while len(boundary)>0:
        round +=1
        next=[]
        for n in boundary:
            for nn in G.neighbors(n):
                if (not nn in boundary+done+next):
                    next.append(nn)
#                    r[nn]=round
        nr={}
        for b in next:
            if len( [nn for nn in G.neighbors(b) if nn not in r] )==1: nr[b]=round
        r.update(nr)
        for b in boundary: done.append(b)
        boundary = []
        for n in next:
            if n in r: boundary.append(n)
#        next=[]
    for n in G.nodes():
        if n not in r: r[n]=round
    return r

def log(x):
    sys.stderr.write(x+"\n")

LOCAL_BRIDGE=1
CUT_ME=2

edge_tags={}

def add_tag(s1,s2,tag):
    if s1<s2:
        ot = edge_tags.get((s1,s2),set())
        ot.add(tag)
        edge_tags[s1,s2] = ot
    else:
        ot = edge_tags.get((s2,s1),set())
        ot.add(tag)
        edge_tags[s2,s1] = ot

def get_tags(s1,s2):
    if s1<s2:
        ot = edge_tags.get((s1,s2),set())
        return tuple(ot)
    else:
        ot = edge_tags.get((s2,s1),set())
        return tuple(ot)

edge_color_setting="hair"
def edge_tag_to_style(tags,setting=edge_color_setting):
    if setting == "hair":
        style=""
        if "hair" in tags:
            style= "color=red"
        elif "longHair" in tags:
            style= "color=orange"
        elif "H" in tags:
            style= "color=blue"
        elif "Y" in tags:
            style= "color=goldenrod"
        elif "nearY" in tags:
            style= "color=goldenrod4"
        elif "bigH" in tags:
            style= "color=green"
        if "promisc" in tags:
            style += " style=dashed"
        return style

def printdot(g,gg0,c,n,ll,bh,annot,trim_level={},post_trim_degree={},tag="bad",yDist={},leafDist={},edgeTags=edge_tags):
#def printdot(g,c,n,ll,bh,annot,tag="bad"):
#    print ll

    import colorsys

    chromosomes={}
    chr_mins={}
    chr_maxs={}
    if bh:
        sb=[]
        for cc in c:
            bhi = bh.get(cc,[0,0,0,0,0,0])
            sb.append( ( bhi[1],old_div((float(bhi[3])+float(bhi[4])),2.0),bhi[2],cc ) )
            chromosomes[bhi[1]]=1
            if chr_mins.get(bhi[1],5.0e9)>min( float(bhi[3]), float(bhi[4]) ):  chr_mins[bhi[1]]=min( float(bhi[3]), float(bhi[4]) )
            if chr_maxs.get(bhi[1],-1.0) <max( float(bhi[3]), float(bhi[4]) ):  chr_maxs[bhi[1]]=max( float(bhi[3]), float(bhi[4]) )
        sb.sort()

# {'19': 46194830.0, '18': 59221558.0, '8': 96645227.0, '4': 18548230.0, 'X': 102465955.0}
# {}

    print("#",chr_mins)
    print("#",chr_maxs)

    nchrs=len(list(chromosomes.keys()))
    i=0
    chr_hue={}
    for ch in list(chromosomes.keys()):
        chr_hue[ch] = old_div(float(i),nchrs)
        i+=1

    gg = nx.subgraph(g,c)
    f=open("%s-%d.txt" % (tag,n), "wt")
    nn=1
    lab0={}
    for x in c:
        p=""
        d=x
        if x[0]=="-":
            p="-"
            d=x[1:]
#        lab0[d]=lab0.get(d,nn)
        lab0[d]=lab0.get( d, float(ll.get(d,0)))
        #print x,d,p,lab0[d]
        nn+=1
    lab={}
    node_fill={}
    for x in c:
        p=""
        d=x
        if x[0]=="-":
            p="-"
            d=x[1:]

        bhi = bh.get(x,False)
        if bhi:
            lab[x]="{:.1f} {}{}\\n{:.2f}-{:.2f}\\n{}".format( old_div(lab0.get(d,0.0),1000), bhi[1],bhi[2],old_div(float(bhi[3]),1.0e6),old_div(float(bhi[4]),1.0e6), x)
#            lab[x]="{:.1f} {}{}\\n{:.2f}-{:.2f}\\n{} {} {} {}".format( lab0.get(d,0.0)/1000, bhi[1],bhi[2],float(bhi[3])/1.0e6,float(bhi[4])/1.0e6,leafDist.get(x,""),yDist.get(x,""), trim_level[x], post_trim_degree.get(x,"") )
            if ( chr_maxs.get(bhi[1], (1.0+float(bhi[3])+float(bhi[4])) ) ) ==0.0: #/2.0)-chr_mins.get(bhi[1],0.0))==0.0:
                print("wtf?",x,bhi,bhi[1],( chr_maxs.get(bhi[1], (1.0+float(bhi[3])+float(bhi[4])) ) ))
            rgb=(0,0,0)
            try:
                rgb=colorsys.hls_to_rgb( chr_hue[bhi[1]], 0.5, old_div((old_div((float(bhi[3])+float(bhi[4])),2.0) - chr_mins.get(bhi[1],0)),(chr_maxs.get(bhi[1],old_div((1.0+float(bhi[3])+float(bhi[4])),2.0))-chr_mins.get(bhi[1],0.0)))  )
            except Exception as e:
                print(e)
            node_fill[x]= '#%02x%02x%02x' % (255.0*rgb[0], 255.0*rgb[1], 255.0*rgb[2] ) #"#{}{}{}".format()
        else:
            lab[x]="{:.1f}\\n{}".format(old_div(lab0.get(d,0.0),1000),str(x))
            node_fill[x]="white"
#        lab[x]="{} {}".format( trim_level.get(x,"?"), post_trim_degree.get(x,"?") )

    f.write( "graph G {\n")
#    f.write( "node [margin=0 fontcolor=blue fontsize=32 width=0.5 shape=circle style=filled]")
    f.write( "node [margin=0 fontsize=6 shape=box];\n")
    f.write( "edge [ fontsize=6 ];\n")

    for x in list(lab.keys()):
        f.write( "{0} [label=\"{1}\" fillcolor=\"{2}\" style=\"filled\" color=\"{2}\"] ; \n".format(x,lab[x],node_fill[x]) )

    if bh:
        last=False
        lastx=0.0
        lastc=0
        for c in sb:
            if last and c[0]==last and (c[1]-lastx)<1000000:
                last_bhi = bh.get(lastc,False)
                this_bhi = bh.get(c[-1],False)
                blast_label=str(last_bhi) + str(this_bhi)
                if this_bhi and last_bhi :
                    aa = tuple(last_bhi[1:5])
                    bb = tuple(this_bhi[1:5])
                    qd = qdist(aa,bb)
                    blast_label = "{}".format(qd)

                if gg0.has_edge(lastc,c[-1]) and not t.has_edge(lastc,c[-1]):
                    f.write("\t \"{}\" -- \"{}\" [weight=2 style=dotted label=\"{} {}\" fontcolor=red] ;\n".format(lastc,c[-1],blast_label,int(abs(gg0[lastc][c[-1]]['weight']))))
                else:
                    f.write("\t \"{}\" -- \"{}\" [weight=2 style=dotted label=\"{}\" fontcolor=blue] ;\n".format(lastc,c[-1],blast_label))
            last=c[0]
            lastx=c[1]
            lastc=c[-1]

    for e in gg.edges():
#        f.write( "\t\"%s\" -- \"%s\";\n" % (lab[e[0]],lab[e[1]]) ) #,gg[e[0]][e[1]]['weight'])
#        color="black"
#        if annot.get(e,0)&LOCAL_BRIDGE : color="red"
#        if annot.get(e,0)&CUT_ME       : color="yellow"
        f.write( "\t \"%s\" -- \"%s\" [label=\"%d\" weight=1 %s];\n" % ( e[0],e[1],int(abs(gg[e[0]][e[1]]['weight'])),edge_tag_to_style( get_tags(e[0],e[1]) ) ))

    f.write( "}\n")

def independent_path(G,a,b,k,t):
    q=[a]
    l={}
    l[a]=0
    r=[]
    while len(q)>0:
#        print a,b,G[a][b],q,r
        n=q.pop(0)
        r.append(n)
        for nn in G.neighbors(n):
            if (n==a and nn==b) or (n==b and nn==a): continue
            if G[n][nn]['weight']>-t: continue
            if nn==b: return True
#            print q,[l[i] for i in q]
            l[nn] = min(l.get(nn,10000), l[n]+1)
            if (not nn in q+r) and (l[nn]<=k):
                q.append(nn)
    return False

def annotate_edges(t,G,node_list):
    an={}
#    nn= len(list(t.edges()))
#    i=0.0
    for a,b in t.edges():

        if not independent_path(G,a,b,4,2):
            an[a,b]="local_bridge"
            an[b,a]="local_bridge"

    return an

def pairs_overlap(x,y):
    a=min(x[0],x[1])
    b=max(x[0],x[1])
    c=min(y[0],y[1])
    d=max(y[0],y[1])

    if a<=c and c<=b: return True
    if a<=d and d<=b: return True
    if c<=a and a<=d: return True
    if c<=b and b<=d: return True
    return False

def qdist(x,y):
    if (not x) or (not y): return (-1)
    if x[0]==y[0]:
        x,y,w,z = list(map(int,[x[2],x[3],y[2],y[3]]))
        ol = pairs_overlap((x,y),(w,z))
        if ol:
            return(-1*min(
                abs(x-w),
                abs(x-z),
                abs(y-w),
                abs(y-z )))
        else:
            return(min(
                abs(x-w),
                abs(x-z),
                abs(y-w),
                abs(y-z )))
    else:
        return(1.0e12)

def test_inversion_option(c1,c2,join_options,ograph,linked):
    a=c1
    b=linked[c1]
    c=c2
    d=linked[c2]

    #k=linked[internal_node]
    ograph.remove_edge(a,b)
    ograph.remove_edge(c,d)
                                                    ##                                                                        ---b a-------c d----
    if a in nx.node_connected_component(ograph,c):  ## exchance labels of a,b if necessary, so the nodes are in this config:  ---a b-------c d------
        x=b
        b=a
        a=x

                                                    ##                                                                        ---b a-------d c----
    if a in nx.node_connected_component(ograph,d):  ## exchance labels of a,b if necessary, so the nodes are in this config:  ---a b-------c d------
        x=b
        b=a
        a=x

        x=d
        d=c
        c=x

                                                    ##                                                                        ---a b-------d c----
    if b in nx.node_connected_component(ograph,d):  ## exchance labels of c,d if necessary, so the nodes are in this config:  ---a b-------c d------
        x=d
        d=c
        c=x

    n_scaffold = old_div(len(nx.node_connected_component(ograph,b)),2)
    print("inversion n nodes",n_scaffold)

    total_i_len = sum( ograph[i][j]['length'] for i,j in nx.bfs_edges(ograph,b) )
    print("inv len",total_i_len)
    if total_i_len < 10000.0 or n_scaffold<2:
        print("inversion length",total_i_len,n_scaffold,"too short")
        join_options.append( (0.0,(),() ) )
        ograph.add_edge(a,b,length=default_gapsize,contig=False)
        ograph.add_edge(c,d,length=default_gapsize,contig=False)

        return

    interc_score0 = intercalation_score_raw(a,b,d,ograph)
    interc_score1 = intercalation_score_raw(a,c,d,ograph)
    print("inversion0", interc_score0)
    print("inversion", interc_score1)
    join_options.append( (interc_score0,(),() ) )
    join_options.append( (interc_score1,((a,c),(b,d)),((a,b),(c,d)) ) )
    ograph.add_edge(a,b,length=default_gapsize,contig=False)
    ograph.add_edge(c,d,length=default_gapsize,contig=False)

    return

def test_endInversion_option(free_end,internal_node,join_options,ograph,linked):
    if free_end in linked:
        x=free_end
        free_end = internal_node
        internal_node = x
    print("end inversion",free_end,linked.get(free_end),internal_node,linked[internal_node])

    k=linked[internal_node]
    ograph.remove_edge(internal_node,k)

    if free_end in nx.node_connected_component(ograph,internal_node):
        x=k
        k=internal_node
        internal_node=x

    sc = link_test(ograph,k,internal_node)
    join_options.append( (sc,(),()))
    print("end inversion existing:",sc)

    sc = link_test(ograph,free_end,internal_node)
    join_options.append( (sc ,((free_end,internal_node),),((internal_node,k),) ) )
    print("end inversion:",sc)

    ograph.add_edge(internal_node,k,length=default_gapsize, contig=False)
    return

def test_interc_option( gap_edge1, gap_edge2, free_end1, join_options, ograph ):
    if not ograph.has_edge(gap_edge1,gap_edge2):
        print("expected nodes to be connected: {} {}".format(gap_edge1, gap_edge2))
        raise Exception('not connected')
    ograph.remove_edge( gap_edge1, gap_edge2 )
    if gap_edge1 in nx.node_connected_component(ograph,gap_edge2):
        print("problem: these should be disconnected now:",gap_edge1,gap_edge2)
        raise Exception('too connected i')
    if gap_edge2 in nx.node_connected_component(ograph,free_end1):
        print("problem: these should have been disconnected all along",free_end1,gap_edge2)
        raise Exception('too connected j')
    if gap_edge1 in nx.node_connected_component(ograph,free_end1):
        print("problem: these should have been disconnected all along",free_end1,gap_edge1)
        raise Exception('too connected k')
    interc_score = intercalation_score(gap_edge1,free_end1,gap_edge2,ograph)
    ograph.add_edge(gap_edge1, gap_edge2, length=default_gapsize,contig=False)
    d = far_end(ograph,free_end1)
    join_options.append( (interc_score,((gap_edge1,free_end1),(d,gap_edge2)),((gap_edge1, gap_edge2),) ) )

N=76710553.0 #100000000.0
pn=0.3
GenomeSize=3.0e9
import math

def get_score(c1_stranded,c2_stranded,gaplen=default_gapsize,cache={}):
    if gaplen<0.0: print(math.log(-1))
#    if cache.has_key((c1_stranded,c2_stranded,gaplen)):
#        return cache[c1_stranded,c2_stranded,gaplen]
#    if cache.has_key((c2_stranded,c1_stranded,gaplen)):
#        return cache[c2_stranded,c1_stranded,gaplen]
    o1,o2=0,0
    if c1_stranded[-2:]==".5": o1=1
    if c2_stranded[-2:]==".3": o2=1
    c1=c1_stranded[:-2]
    c2=c2_stranded[:-2]
    l1=ll[c1]
    l2=ll[c2]
    p0 = ces.p_not_a_hit(l1,l2,GenomeSize,gaplen,pn)
    if p0<0.0:
        print("wtf?  p0<0",p0,l1,l2,GenomeSize,gaplen,pn)
    thisscore=ces.llr_v0(l1,l2,o1,o2,GenomeSize,pn,links.get((c1,c2),[]),N,gaplen,p0 )
    #cache[c1_stranded,c2_stranded,gaplen]=thisscore
    return thisscore

def traverse_and_layout(n,coords,facing,x,s,og,maxD=2000000):
    color={}
    q=[n]
    ptr=x
    if s>0:
        facing[n]="L"
    else:
        facing[n]="R"
    while (abs(ptr-x)<maxD) and len(q)>0:
        m=q.pop(0)
        coords[m]=ptr
        color[m]=1
        for mm in og.neighbors(m):
            if not mm in color:
                q.append(mm)
                ptr+=s*og[m][mm]['length']
                if not og[m][mm]['contig']:
                   if s==-1:
                       facing[mm]='R'
                   else:
                       facing[mm]='L'
                else:
                   if s==-1:
                       facing[mm]='L'
                   else:
                       facing[mm]='R'

def far_end(g,n):
    if not g.degree(n)==1:
        print("wtf: this should be a leaf")
        exit(0)
    for m in nx.node_connected_component(g,n):
        if (not m==n) and g.degree(m)==1:
            return m

def both_ends(g,n):
    p=[]
    for m in nx.node_connected_component(g,n):
        if  g.degree(m)==1:
            p.append( m)
    return p

def strand_check(a,b,x,f):
    if x[a]<x[b]:
        if f[a]=="R" and f[b]=="L": return True
    else:
        if f[b]=="R" and f[a]=="R": return True
    return False

def take_best_join_option(join_options,og,linked,edge_options,label=""):
#    print join_options
    join_options.sort(reverse=True)
#    print join_options
    if join_options[0][0]>0.0:
#        if len(join_options[0][2]) >=1 : print "#accept join",label, join_options[0], join_options
        for x,y in join_options[0][2]:
            og.remove_edge(x,y)
            del linked[x]
            del linked[y]
        for x,y in join_options[0][1]:
            linked[x]=y
            linked[y]=x
            og.add_edge(x,y,length=default_gapsize,contig=False)
            edge_options[x,y] = "[color=brown label=\"{:.2f}\"]".format(weights.get((x,y),0))
            edge_options[y,x] = "[color=brown label=\"{:.2f}\"]".format(weights.get((x,y),0))

def link_test(og,a,b,gapsize=default_gapsize,max_stretch=200000):
    coords = {}
    facing={}
    traverse_and_layout(a,coords,facing,0,-1,og,maxD=200000)
    traverse_and_layout(b,coords,facing,gapsize,+1,og,maxD=200000)
#    print "#x:setup done"
    sys.stdout.flush()

    score=0.0
    for n1 in nx.node_connected_component(og,a):
        if abs(coords.get(n1,-1e10))>max_stretch: continue
        for n2 in nx.node_connected_component(og,b):
            if abs(coords.get(n2,1e10))>max_stretch: continue
            if strand_check(n1,n2,coords,facing):
                distance = coords[n2] - coords[n1]
                if distance < max_stretch:
                    x=get_score(n1,n2,distance)
#                    print "#x:partial",n1,n2,x
                    sys.stdout.flush()
                    score += x
    return score

def intercalation_score(a,b,c,og):
#    print "#intercalation test:",a,b,c
    if a in nx.node_connected_component(og,b):
        print("a should not be connected to b",a,b)
        raise Exception('too connected a')
    if b in nx.node_connected_component(og,c):
        print("b should not be connected to c",b,c)
        raise Exception('too connected b')
    if a in nx.node_connected_component(og,c):
        print("a should not be connected to c",a,c)
        raise Exception('too connected c')
    coordinates1={}
    coordinates2={}
    facing={}
    traverse_and_layout(a,coordinates1,facing,0,-1,og)
    coordinates2=dict(coordinates1)
    traverse_and_layout(c,coordinates1,facing,default_gapsize,+1,og)
    traverse_and_layout(b,coordinates2,facing,default_gapsize,+1,og)
    traverse_and_layout(c,coordinates2,facing,default_gapsize+max(coordinates2.values()),+1,og)

    score_0=0.0
    for n1 in nx.node_connected_component(og,a):
        for n2 in nx.node_connected_component(og,c):
            if strand_check(n1,n2,coordinates1,facing):
                distance = coordinates1[n2] - coordinates1[n1]
                if distance < 200000:
                    if distance<0: print("wtf1?")
                    score_0 += get_score(n1,n2,distance)

    score_1=0.0
    for n1 in nx.node_connected_component(og,a):
        for n2 in nx.node_connected_component(og,c):
            if strand_check(n1,n2,coordinates2,facing):
                distance = coordinates2[n2] - coordinates2[n1]
                if distance < 200000:
                    if distance<0: print("wtf2?")
                    score_1 += get_score(n1,n2,distance)

        for n2 in nx.node_connected_component(og,b):
            if strand_check(n1,n2,coordinates2,facing):
                distance = coordinates2[n2] - coordinates2[n1]
                if distance < 200000:
                    if distance<0: print("wtf3?")
                    score_1 += get_score(n1,n2,distance)

    for n1 in nx.node_connected_component(og,b):
        for n2 in nx.node_connected_component(og,c):
            if strand_check(n1,n2,coordinates2,facing):
                distance = coordinates2[n2] - coordinates2[n1]
                if distance < 200000:
                    if distance<0: print("wtf4?")
                    score_1 += get_score(n1,n2,distance)

#    print "#intercalation scores:",score_1,score_0

    if a in nx.node_connected_component(og,b):
        print("a should not be connected to b",a,b)
        raise Exception('too connected d')
    if b in nx.node_connected_component(og,c):
        print("b should not be connected to c",b,c)
        raise Exception('too connected e')
    if a in nx.node_connected_component(og,c):
        print("a should not be connected to c",a,c)
        raise Exception('too connected f')

    return (score_1-score_0)

def intercalation_score_raw(a,b,c,og):
#    print "#intercalation test:",a,b,c
    if a in nx.node_connected_component(og,b): print("a should not be connected to b",a,b)
    if b in nx.node_connected_component(og,c): print("b should not be connected to c",b,c)
    if a in nx.node_connected_component(og,c): print("a should not be connected to c",a,c)
#    coordinates1={}
    coordinates2={}
    facing={}
    traverse_and_layout(a,coordinates2,facing,0,-1,og)
#    coordinates2=dict(coordinates1)
#    traverse_and_layout(c,coordinates1,facing,1000,+1,og)
    traverse_and_layout(b,coordinates2,facing,default_gapsize,+1,og)
    traverse_and_layout(c,coordinates2,facing,default_gapsize+max(coordinates2.values()),+1,og)

#    score_0=0.0
#    for n1 in nx.node_connected_component(og,a):
#        for n2 in nx.node_connected_component(og,c):
#            if strand_check(n1,n2,coordinates1,facing):
#                distance = coordinates1[n2] - coordinates1[n1]
#                if distance < 200000:
#                    if distance<0: print "wtf1?"
#                    score_0 += get_score(n1,n2,distance)

    score_1=0.0
    for n1 in nx.node_connected_component(og,a):
        for n2 in nx.node_connected_component(og,c):
            if strand_check(n1,n2,coordinates2,facing):
                distance = coordinates2[n2] - coordinates2[n1]
                if distance < 200000:
                    if distance<0: print("wtf2?")
                    score_1 += get_score(n1,n2,distance)

        for n2 in nx.node_connected_component(og,b):
            if strand_check(n1,n2,coordinates2,facing):
                distance = coordinates2[n2] - coordinates2[n1]
                if distance < 200000:
                    if distance<0: print("wtf3?")
                    score_1 += get_score(n1,n2,distance)

    for n1 in nx.node_connected_component(og,b):
        for n2 in nx.node_connected_component(og,c):
            if strand_check(n1,n2,coordinates2,facing):
                distance = coordinates2[n2] - coordinates2[n1]
                if distance < 200000:
                    if distance<0: print("wtf4?")
                    score_1 += get_score(n1,n2,distance)

#    print "#intercalation scores:",score_1,score_0
    return (score_1)

if __name__=="__main__":

    import sys
    import argparse
    parser = argparse.ArgumentParser()

    parser.add_argument('-t','--threshold',default=0.0 ,  type=float)

    parser.add_argument('-H','--head',default=False,type=int)
    parser.add_argument('-D','--savetreedots',default=False,action='store_true')
    parser.add_argument('-d','--debug',default=False,action='store_true')
#    parser.add_argument('-I','--nointerc',default=False,action='store_true')
    parser.add_argument('-p','--progress',default=False,action='store_true')
    parser.add_argument('-M','--maxdegree',default=False,type=int)
    parser.add_argument('-m','--minlength',default=500,type=int)

    parser.add_argument('-S','--silent',default=False,action='store_true')
    parser.add_argument('-K','--cutPromisc',default=False,action='store_true')
    parser.add_argument('-J','--logH',default=False,action='store_true')
    parser.add_argument('-T','--logTags',default=False,action='store_true')
    parser.add_argument('-C','--cheat',default=False,action='store_true')
    parser.add_argument('-B','--blacklist')

    parser.add_argument('-j','--joins') #pre-join these
    parser.add_argument('-k','--links')
    parser.add_argument('-b','--besthits')
    parser.add_argument('-s','--skip',default=0,type=int)
    parser.add_argument('-l','--lengths')
    parser.add_argument('-E','--edgefile')
    parser.add_argument('--set_insert_size_dist_fit_params')
    parser.add_argument('-L','--maxLength',type=float,default=150000.0)
    parser.add_argument('-P','--promisc',type=float,default=0.023)
    parser.add_argument('-o','--dotLabel',default="bad")

    args = parser.parse_args()
    if args.debug:
        args.progress=True

    if args.progress: log( str(args) )
    print("#"+str(args))

    if args.set_insert_size_dist_fit_params:
        s=args.set_insert_size_dist_fit_params
        a,b,c,d,f,pn,N = list(map(float,list(s.split(','))))
        print("#",a,b,c,d,f)
        ces.set_insert_size_dist_fit_params(a,b,c,d,f,pn,N)
        sys.stdout.flush()

    G=nx.Graph()
    SG=nx.Graph()

    ll={}
    if args.lengths:
        f = open(args.lengths)
        while True:
            l = f.readline()
            if not l: break
            if l[0]=="#": continue

            c=l.strip().split()
            l = int(c[1])
            ll[c[0]]=int(c[1])
            if l>= args.minlength:
                G.add_node(c[0])
                SG.add_node(c[0])
        f.close()
    if args.progress: print("#Done reading lengths")

    #broken.al.masked.links.txt
    def read_links(contigs=False):
        links={}
        import glob
        for f in glob.glob(args.links): #["same_component-t4.links.txt"]:
            print("#file:",f)
            ff=open(f)
            while True:
                l = ff.readline()
                if not l: break
                if l[0]=="#": continue
                c=l.strip().split("\t")
                if (not contigs) or (c[0] in contigs) and (c[1] in contigs):
                    links[(c[0],c[1])]=eval(c[5])
        return links

    besthit={}
    if args.besthits:
#        besthit={}
        if args.besthits:
            f = open(args.besthits)
            while True:
                l = f.readline()
                if not l: break

                if not l[:5]=="best:": continue
                c=l.strip().split()
                besthit[c[1]]=c[2:]
    #            print c[1],besthit[c[1]]
            f.close()
    if args.progress: print("#Done reading besthits")

    if args.edgefile:
        f = open(args.edgefile)
    else:
        f=sys.stdin
    while True:
        l = f.readline()
        if not l: break
        if l[0]=="#": continue
        c=l.strip().split()
        u,v,w = c[0],c[1],float(c[2])
        if ( not args.lengths ) or (ll[u]>=args.minlength and ll[v]>=args.minlength):
            G.add_edge(u,v,weight=-w)
            if w >= args.threshold:
                SG.add_edge(u,v,weight=-w)
    if args.edgefile:
        f.close()

    if args.progress: print("#Done reading edgelist")

    bad_nodes=[]
    total_discarded_length=0
    total_discarded_length1=0
    n_discarded1=0
    n_discarded2=0
    if args.maxdegree:
        for n in SG.nodes():
            print("#dg:", SG.degree(n))
            if SG.degree(n)>args.maxdegree:
                n_discarded1+=1
                bad_nodes.append(n)
                total_discarded_length += ll[n]
                print("#discard:",n,ll[n],SG.degree(n))
                for nn in SG.neighbors(n):
                    if SG.degree(nn)==1:
                        n_discarded2+=1
                        total_discarded_length1+=ll[nn]
        for n in bad_nodes:
            e_to_remove=[]
            for e in SG.edges([n]):
                e_to_remove.append(e)
            SG.remove_edges_from(e_to_remove)

    if args.cutPromisc:
        bad_nodes=[]
        for n in G.nodes():
#            print "#dg:", SG.degree(n)
            if (old_div(float(G.degree(n)), ll[n]))>args.promisc : # G.degree(n)/ll[n]>args.maxdegree:
                n_discarded1+=1
                bad_nodes.append(n)
                total_discarded_length += ll[n]
                print("#discard:",n,ll[n],G.degree(n))
#                for nn in G.neighbors(n):
#                    if G.degree(nn)==1:
#                        n_discarded2+=1
#                        total_discarded_length1+=ll[nn]
        for n in bad_nodes:
            e_to_remove=[]
            for e in G.edges([n]):
                e_to_remove.append(e)
            G.remove_edges_from(e_to_remove)
            SG.remove_edges_from(e_to_remove)

    if args.blacklist:
        f=open(args.blacklist)
        e_to_remove=[]
        while True:
            l=f.readline()
            if not l: break
            c=l.strip().split()
            e_to_remove.append((c[1],c[2]))
        G.remove_edges_from(e_to_remove)
        SG.remove_edges_from(e_to_remove)

    if args.cheat:
        e_to_remove=[]
        for a,b in SG.edges():
            if not ( a in besthit and b in besthit):
                e_to_remove.append((a,b))
            else:
                aa = tuple(besthit[a][1:5])
                bb = tuple(besthit[b][1:5])
                qd = qdist(aa,bb)
                if qd >= args.maxLength :
                    e_to_remove.append((a,b))
        SG.remove_edges_from(e_to_remove)

    if args.progress: print("#total_discarded_length",n_discarded1,n_discarded2,old_div(float(total_discarded_length),1.0e6),old_div(float(total_discarded_length1),1.0e6),old_div(float(total_discarded_length+total_discarded_length1),1.0e6))

    promisc = {}
    for n in G.nodes():
        if args.debug: print("#r:",old_div(float(G.degree(n)),ll[n]))
        if (old_div(float(G.degree(n)),ll[n]))>args.promisc: promisc[n]=True

    tag_tallies={}
    bad_tag_tallies={}
    strx={"+":0, "-":1}
    strings = []
    ccn=1

    if args.progress:
        print("about to load the links")
        sys.stdout.flush()
    links = read_links()
    if args.progress:
        print("loaded the links")
        sys.stdout.flush()

    print("#n_connected_components:",len(nx.connected_components(SG)))
    for c in nx.connected_components(SG):
        if ccn<=args.skip:
            ccn+=1
            continue

        if len(c)==1:
            print("#edge:",c[0]+".5",c[0]+".3",{"length": ll[c[0]],"contig": True})
            continue

#        if args.progress:
#            print "about to load the links"
#            sys.stdout.flush()
#        links = read_links(c)
#        if args.progress:
#            print "loaded the links"
#            sys.stdout.flush()

        print("cc:",ccn,len(c),c)

        edge_options={}
        nodes={}

        og=nx.Graph()
        linked={}

        for cc in c:
            og.add_node(cc+".5")
            og.add_node(cc+".3")
            nodes[cc+".5"]="5'"
            nodes[cc+".3"]="3'"
            og.add_edge(cc+".5",cc+".3",length=ll[cc],contig=True)
            edge_options[cc+".5",cc+".3"] = "[label=\"{} {}\" penwidth=3]".format(cc,old_div((ll[cc]+500),1000))
            edge_options[cc+".3",cc+".5"] = "[label=\"{} {}\" penwidth=3]".format(cc,old_div((ll[cc]+500),1000))

        if args.joins:
            f = open(args.joins)
            while True:
                l = f.readline()
                if not l: break
                cz=l.strip().split()
                if not len(cz)>=3: continue
                #print "z:",cz[1][:-2],cz[2][:-2],cz
                if cz[0]=="#join:" :
                    if (cz[1][:-2] in c) and (cz[2][:-2] in c):
                        print("join",cz[1:])
                        og.add_edge( cz[1],cz[2],length=default_gapsize,contig=False)
                        linked[cz[1]]=cz[2]
                        linked[cz[2]]=cz[1]
#                        og.add_edge(c1,c2,length=1000.0,contig=False)
                        edge_options[cz[1],cz[2]] = "[label=\"pre\" penwidth=2]"
                        edge_options[cz[2],cz[1]] = "[label=\"pre\" penwidth=2]"
                    else:
                        pass
#                        print "skip",cz[1:]
            f.close()

        weights = {}
        for c1 in c:
            if args.progress: print("#",c1)
            l1=ll[c1]
            for c2 in c:
                if not c1<c2: continue
                if args.joins and c2+".3" in nx.node_connected_component(og, c1+".3"): continue
                l2=ll[c2]
                if (c1,c2) not in links: continue
                for gaplen in [default_gapsize]: # [ 0, 500, 1000, 2000 , 10000, 20000, 30000, 40000, 100000, 200000, 500000 ]:
                    p0 = ces.p_not_a_hit(l1,l2,GenomeSize,gaplen,pn)
                    s={}
                    for (o1,o2,suf1,suf2) in ((0,0,".3",".5"),(0,1,".3",".3"),(1,0,".5",".5"),(1,1,".5",".3")):
                        weights[c1+suf1,c2+suf2] = ces.llr_v0( l1,l2,o1,o2,GenomeSize,pn,links.get((c1,c2),[]),N,gaplen,p0 )
                        weights[c2+suf2,c1+suf1] = weights[c1+suf1,c2+suf2]
                        if args.progress: print("#",c1+suf1,c2+suf2,len(links[c1,c2]),weights[c1+suf1,c2+suf2])

        link_pairs = list(weights.keys())
        link_pairs.sort(key = lambda x: weights[x], reverse=True)

        for c1,c2 in link_pairs:
            if not c1<c2: continue
            if weights[c1,c2]>12:
                if c2 in nx.node_connected_component(og, c1):  # link within one of the scaffolds.
                    if c1 in linked and linked[c1]==c2: continue # this join pre made!
                    if (not c1 in linked) and (not c2 in linked):  # this would circularize
                        pass
                    elif ( c1 in linked) and ( c2 in linked):      # test to invert (or excize circle?)
                        if og.has_edge(c1,c2):
                            print("skip test flip of",c1,c2)
                            continue #don't test inversions of individual contigs
#                        print "test invert",c1,linked.get(c1),c2, linked.get(c2)
                        join_options=[]
                        test_inversion_option(c1,c2,join_options,og,linked)
                        take_best_join_option(join_options,og,linked,edge_options,"inversion")

                    elif (not c1 in linked) or (not c2 in linked): # test for end inversion (scorpion tail?) or pinch off circle from end?
#                        print "test end invert",c1,c2,linked.get(c1), linked.get(c2)
                        join_options=[]
                        test_endInversion_option(c1,c2,join_options,og,linked)
                        take_best_join_option(join_options,og,linked,edge_options,"end inversion")

                else:
                    if (not c1 in linked) and (not c2 in linked): # and not c2 in nx.node_connected_component(og, c1):
                        linked[c1]=c2
                        linked[c2]=c1
                        print("#easy",c1,c2,weights[c1,c2])
                        og.add_edge(c1,c2,length=default_gapsize,contig=False)
                        edge_options[c1,c2] = "[color=blue label=\"{:.2f}\"]".format(weights[c1,c2])
                        edge_options[c2,c1] = "[color=blue label=\"{:.2f}\"]".format(weights[c1,c2])
                    elif (c1 in linked) and (c2 in linked):
                        # This edge looks like an H.  test whether the two linked scaffolds can be put end-to-end:
                        # a-b k-d : we'll test b+k, a+k, b+d and a+d, and if the best one raises the total score, we'll add it.
                        a,b = both_ends(og,c1)
                        k,d = both_ends(og,c2)
                        sc={}
                        for x,y in ((b,k),(a,k),(b,d),(a,d)):
                            sc[x,y] = link_test(og,x,y)
                        sll=list(sc.items())
                        sll.sort(key=lambda x: x[1],reverse=True)
                        print("# Htest:", c1,c2,sll)
                        x,y = sll[0][0]
                        if sc[x,y]>0:
                            print("#H accept", sc[x,y] ,x,y)
                            linked[x]=y
                            linked[y]=x
                            og.add_edge(x,y,length=default_gapsize,contig=False)
                            edge_options[x,y] = "[color=green label=\"{:.2f}\"]".format(weights.get((x,y),0))
                            edge_options[y,x] = "[color=green label=\"{:.2f}\"]".format(weights.get((x,y),0))

                    elif (not c1 in linked) or (not c2 in linked):  #test whether one can be intercalated in the other:
                        accepted=False
                        join_options=[]
                        if (not c1 in linked):
                            b=c1
                            a=c2
                            cP=linked[a]
                            print("#1:",c1,c2,a,b,cP,weights[c1,c2])
                        elif (not c2 in linked):
                            b=c2
                            a=c1
                            cP=linked[a]
                            print("#2:",c1,c2,a,b,cP,weights[c1,c2])
                        if (not c1 in linked) or (not c2 in linked):
                            # not linking end-to-end:
                            #but one of them IS an end link, try inserting that scaffold into the other one.
                            # a and c are the contig ends flanking the gap we're inserting into, b is the end of
                           #  the scaffold being tested for insertion.

                            test_interc_option(a,cP,b,join_options,og)
                            test_interc_option(a,cP,far_end(og,b),join_options,og)

                            a,k = both_ends(og,a)
                            b,d = both_ends(og,b)
                            sc={}
                            for x,y in ((a,b),(k,b),(a,d),(k,d)):
                                sc[x,y] = link_test(og,x,y)
                                join_options.append( (sc[x,y],((x,y),),() ) )

                            take_best_join_option(join_options,og,linked,edge_options,"interc or end")

        if False:
            for c1,c2 in link_pairs:
                if weights[c1,c2]>12 and not c2 in nx.node_connected_component(og, c1):
    #                linked[c1]=True
    #                linked[c2]=True
                    og.add_edge(c1,c2)
                    edge_options[c1,c2] = "[style=dotted label=\"{:.2f}\"]".format(weights[c1,c2]) #]"
                    edge_options[c2,c1] = "[style=dotted label=\"{:.2f}\"]".format(weights[c1,c2]) #]"

        red_edges=""
        node_fill={}
        if args.besthits:
            import colorsys

            chromosomes={}
            chr_mins={}
            chr_maxs={}
            ends=[]

            for cc in c:
#best:   Scaffold91_1    46380.0 chr12   +       49064574        49095652
#                besthit[c[1]]=c[2:]

                if cc not in besthit: continue
                aa = tuple(besthit[cc][1:5])
                chromosomes[aa[0]]=1
                xxx = old_div((int(aa[2])+int(aa[3])),2)
                if chr_mins.get(aa[0],1.0e9)>xxx: chr_mins[aa[0]]=xxx
                if chr_maxs.get(aa[0],-1.0e9)<xxx: chr_maxs[aa[0]]=xxx
                if aa[1]=="+":
                    ends.append((aa[0],xxx,cc+".5",cc+".3",cc))
#                    ends.append((c[0],int(aa[3]),cc+".3",cc))
                elif aa[1]=="-":
                    ends.append((aa[0],xxx,cc+".3",cc+".5",cc))
#                    ends.append((c[0],int(aa[3]),cc+".5",cc))
            ends.sort()

            nchrs=len(list(chromosomes.keys()))
            if nchrs==0: nchrs+=1
            i=0
            chr_hue={}
            for ch in list(chromosomes.keys()):
                chr_hue[ch] = old_div(float(i),nchrs)
                i+=1

            for cc in c:
                if not cc in besthit: continue
                aa = tuple(besthit[cc][1:5])
                chrid=aa[0]
                xxx = old_div(float(int(aa[2])+int(aa[3])),2)
                rgb=colorsys.hls_to_rgb( chr_hue[chrid], 0.5, old_div((xxx - chr_mins.get(chrid,0)),(1.0 + chr_maxs.get(chrid,xxx+1.0)-chr_mins.get(chrid,0.0))))
                node_fill[cc+".5"]= '#%02x%02x%02x' % (255.0*rgb[0], 255.0*rgb[1], 255.0*rgb[2] ) #"#{}{}{}".format()
                node_fill[cc+".3"]= '#%02x%02x%02x' % (255.0*rgb[0], 255.0*rgb[1], 255.0*rgb[2] ) #"#{}{}{}".format()

            for i in range(1,len(ends)):
                if ends[i-1][0]==ends[i][0] and ends[i-1][3] in nx.node_connected_component(og,ends[i][3]):
#                    red_edges += "\"{}\" -- \"{}\" [color=red style=dotted label=\"{:.1f}\" fontcolor=red];\n".format(ends[i-1][3],ends[i][2],weights.get((ends[i-1][3],ends[i][2]),""))
                    red_edges += "\"{}\" -- \"{}\" [color=red style=dotted ];\n".format(ends[i-1][3],ends[i][2])

        if args.savetreedots:
            f=open("greedy-linear7-{}.dot".format(ccn),"wt")
            f.write("graph G {\n")
            for cc in list(nodes.keys()):
    #[label=\"{1}\" fillcolor=\"{2}\" style=\"filled\" color=\"{2}\"]
                f.write("\"{}\" [ label=\"{}\" fillcolor=\"{}\" color=\"{}\" style=\"filled\" ];\n".format(cc,nodes[cc],node_fill.get(cc,"white"),node_fill.get(cc,"white")))
                #            f.write("\n".format(c))
            f.write(red_edges)
            for c1,c2 in og.edges():
                f.write("\"{}\" -- \"{}\" {} \n".format(c1,c2,edge_options.get((c1,c2),"")))
            f.write("}\n")
            f.close()

        for nnog in nx.connected_component_subgraphs(og):
            tlnnog=0
            for c1,c2 in nnog.edges():
                print("#edge:",c1,c2,og.get_edge_data(c1,c2))
                tlnnog+=og[c1][c2]['length']
            print("#slen",tlnnog)
        ccn+=1

        if args.head and ccn>args.head:
            break


"""Plot figures from the paper."""
import logging
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
from astropy import units as u
from astropy.table import Table
from gammapy.spectrum import CrabSpectrum, CountsPredictor
from gammapy.spectrum.models import LogParabola
import matplotlib.lines as mlines
from matplotlib import gridspec
from .models import Log10Parabola
from .utils import load_yaml
from .conf import config
from .errors import stat_errorband

log = logging.getLogger(__name__)

FONTSIZE = 15
FONTSIZE_CONTOURS = 18
E_UNIT_LABEL = r"$E\,/\,\mathrm{TeV}$"
SED_UNIT_LABEL = (
    r"$E^2 \cdot {\rm d}\phi/{\rm d}E\,/\,({\rm erg}\,{\rm cm}^{-2} {\rm s}^{-1})$"
)
COLORS = ["#21ABCD", "#FF9933", "#893F45", "#3EB489", "#002E63", "crimson"]

def plot_crab():
    """Plot Crab pulsar and nebula SED."""
    log.info("Executing plot_crab ...")

    fig, ax = plt.subplots()

    # Plot flux points
    for component in ["pulsar", "nebula"]:
        table = Table.read("data/other/crab_mwl.fits.gz")
        table = table[table["component"] == component]
        x = table["energy"].data
        y = table["energy_flux"].data
        yerr_lo = table["energy_flux_err_lo"].data
        yerr_hi = table["energy_flux_err_hi"].data
        ax.errorbar(x, y, yerr=(yerr_lo, yerr_hi), fmt="o", label=component)

    # Plot SED model
    energy = np.logspace(2, 8, 100) * u.MeV

    crab = CrabSpectrum(reference="meyer")

    flux = crab.model(energy)
    energy_flux = (energy ** 2 * flux).to("erg cm^-2 s^-1")
    ax.plot(energy.value, energy_flux.value, label="Meyer (2010) model", lw=3)

    ax.set_xlim((3e-1, 3e8))
    ax.set_ylim((3e-12, 3e-8))
    ax.set_xlabel("Energy (MeV)")
    ax.set_ylabel("E^2 dN/dE (erg cm^-2 s^-1)")
    fig.legend(loc="upper center", ncol=3)
    ax.grid()
    ax.loglog()

    path = Path("results/figures")
    path.mkdir(parents=True, exist_ok=True)
    filename = "results/figures/crab_mwl.png"
    log.info(f"Writing {filename}")
    fig.savefig(filename)

def plot_fit_results(tool):
    """plot the SEDs result of the gammapy / sherpa fit
    for comparison with the literature we choose only the Mayer spectrum

    Here we plot the butterfly as a result of the multivariate sampling
    with the 68% containment in flux
    """
    fig, ax = plt.subplots()

    model_meyer_ref = CrabSpectrum("meyer").model
    model_meyer_ref.plot(
        [10 * u.GeV, 100 * u.TeV],
        energy_power=2,
        flux_unit="erg-1 cm-2 s-1",
        ls=":",
        lw=2.2,
        color="#555555",
        label="Meyer et al. (2010)",
    )

    # where to take the results, configurations for the individual butterflies
    instruments = ["fermi", "magic", "veritas", "fact", "hess", "joint"]
    labels = ["Fermi-LAT", "MAGIC", "VERITAS", "FACT", "H.E.S.S.", "joint fit"]
    lss = ["--", "--", "--", "--", "--", "-"]
    colors = COLORS
    # with one loop we realize all the butterfly plots
    for instrument, label, color, ls in zip(instruments, labels, colors, lss):

        path = (
            config.repo_path
            / f"results/fit/{tool}/{instrument}/fit_results_logparabola.yaml"
        )

        if not path.exists():
            log.warning(f"Missing: {path} . Skipping.")
            continue

        results = load_yaml(path)
        parameters = results["parameters"]

        model_lp = LogParabola.from_log10(
            amplitude=parameters[0]["value"] * u.Unit(parameters[0]["unit"]),
            reference=parameters[1]["value"] * u.Unit(parameters[1]["unit"]),
            alpha=parameters[2]["value"] * u.Unit(parameters[2]["unit"]),
            beta=parameters[3]["value"] * u.Unit(parameters[3]["unit"]),
        )

        # energy range for the plot
        dataset = config.get_dataset(instrument)
        energy_range = dataset.energy_range

        # just in case of the joint fit put a thicker line and a less transparent butterfly
        if instrument == "joint":
            model_lp.plot(
                energy_range,
                energy_power=2,
                flux_unit="erg-1 cm-2 s-1",
                ls=ls,
                lw=3,
                color=color,
                label=label,
            )
        else:
            model_lp.plot(
                energy_range,
                energy_power=2,
                flux_unit="erg-1 cm-2 s-1",
                ls=ls,
                lw=2.2,
                color=color,
                label=label,
            )

        # read the butterfly from the multivariate sampling results
        table_path = Path(
            f"{config.repo_path}/results/figures/stat_err/{instrument}_flux_errorband.dat"
        )
        log.info(f"reading butterfly values from {table_path}")
        t = Table.read(table_path, format="ascii.ecsv")
        energies = t["energies"].data * t["energies"].unit
        flux_lo = t["flux_lo"].data * t["flux_lo"].unit
        flux_hi = t["flux_hi"].data * t["flux_hi"].unit

        if instrument == "joint":
            alpha = 0.38
        else:
            alpha = 0.28

        plt.fill_between(
            energies.to("TeV"),
            (energies ** 2 * flux_lo).to("erg cm-2 s-1"),
            (energies ** 2 * flux_hi).to("erg cm-2 s-1"),
            color=color,
            alpha=alpha,
            label="",
        )

    ax.legend(fontsize=FONTSIZE)
    ax.set_ylim([1e-12, 2e-10])

    ax.set_xlabel(E_UNIT_LABEL, size=FONTSIZE)
    ax.set_ylabel(SED_UNIT_LABEL, size=FONTSIZE)
    # make axis thicker
    for axis in ["top", "bottom", "left", "right"]:
        ax.spines[axis].set_linewidth(1.6)
    ax.tick_params("both", length=7, width=1.6, which="major", labelsize=FONTSIZE)
    ax.tick_params("both", length=4, width=1.6, which="minor", labelsize=FONTSIZE)

    plt.tight_layout()

    filename = f"results/figures/crab_sed_{tool}_fit.png"
    filename_pdf = f"results/figures/crab_sed_{tool}_fit.pdf"
    log.info(f"Writing {filename}")
    fig.savefig(filename)
    fig.savefig(filename_pdf)

def plot_sherpa_contours():
    """plot the confidence contours obtained from the sherpa fit
    """
    log.info("plotting parameters contours obtained from sherpa")
    # where to take the results, configurations for the individual butterflies
    instruments = ["fermi", "magic", "hess", "fact", "veritas", "joint"]
    labels = ["Fermi-LAT", "MAGIC", "H.E.S.S.", "FACT", "VERITAS", "joint fit"]
    colors = ["#21ABCD", "#FF9933", "#5A4FCF", "#5CC184", "#702963", "crimson"]
    lss = ["--", "--", "--", "--", "--", "-"]

    fig, axarr = plt.subplots(1, 3, figsize=(18, 6))

    # with one loop we realize all the contour plots
    for instrument, label, color, ls in zip(instruments, labels, colors, lss):

        path = config.repo_path / f"results/fit/sherpa/{instrument}"

        contours_path = path / "fit_contours_logparabola.npy"
        results_path = path / "fit_results_logparabola.yaml"

        if not path.exists():
            log.warning(f"Missing: {path} . Skipping.")
            continue

        # load the contours and the results of the fit
        contours = np.load(contours_path).tolist()
        results = load_yaml(results_path)

        # define a 2 x 2 matrix to visualise the plot
        # we will delete one of the subplots and make something like a corner plot
        # useful variables for the plot
        ampl_range = contours["contour_ampl_c1"]["x0_range"]
        c1_range = contours["contour_ampl_c1"]["x1_range"]
        c2_range = contours["contour_ampl_c2"]["x1_range"]
        # actual values output of the fit
        # remember in sherpa notation: (amplitude->ampl, alpha->c1, beta->c2)
        ampl = results["parameters"][0]["value"]
        c1 = results["parameters"][2]["value"]
        c2 = results["parameters"][3]["value"]

        # axarr[0,0]
        extent = [ampl_range[0] * 1e9, ampl_range[1] * 1e9, c1_range[0], c1_range[1]]

        axarr[0].contour(
            contours["contour_ampl_c1"]["like_values"],
            contours["contour_ampl_c1"]["levels"],
            origin="lower",
            extent=extent,
            colors=color,
            linewidths=(2., 1.5, 1.3),
            linestyles=("-", "--", ":"),
        )

        # print actual value
        axarr[0].plot(ampl, c1, marker="X", markersize=7, color=color)
        axarr[0].set_xlabel(
            r"$f_0 / (\mathrm{TeV} \, \mathrm{cm}^{-2} \mathrm{s}^{-1})$"
        )
        axarr[0].set_ylabel(r"$\Gamma$")

        extent = [ampl_range[0] * 1e9, ampl_range[1] * 1e9, c2_range[0], c2_range[1]]

        axarr[1].contour(
            contours["contour_ampl_c2"]["like_values"],
            contours["contour_ampl_c2"]["levels"],
            origin="lower",
            extent=extent,
            colors=color,
            linewidths=(2., 1.5, 1.3),
            linestyles=("-", "--", ":"),
        )

        # print actual value
        axarr[1].plot(ampl, c2, marker="X", markersize=7, color=color)
        axarr[1].set_ylabel(r"$\beta$")
        axarr[1].set_xlabel(
            r"$f_0 / (\mathrm{TeV} \, \mathrm{cm}^{-2} \, \mathrm{s}^{-1})$"
        )

        extent = [c1_range[0], c1_range[1], c2_range[0], c2_range[1]]

        axarr[2].contour(
            contours["contour_c1_c2"]["like_values"],
            contours["contour_c1_c2"]["levels"],
            origin="lower",
            extent=extent,
            colors=color,
            linewidths=(2., 1.5, 1.3),
            linestyles=("-", "--", ":"),
        )

        # print actual value
        axarr[2].plot(c1, c2, marker="X", markersize=7, color=color)
        axarr[2].set_ylabel(r"$\beta$")
        axarr[2].set_xlabel(r"$\Gamma$")

    # axarr[0,1] is for the legend
    import matplotlib.lines as mlines

    sigma_1 = mlines.Line2D(
        [], [], color="k", marker="", ls="-", lw=2., label=r"1 $\sigma$ contour"
    )
    sigma_2 = mlines.Line2D(
        [], [], color="k", marker="", ls="--", lw=1.5, label=r"2 $\sigma$ contour"
    )
    sigma_3 = mlines.Line2D(
        [], [], color="k", marker="", ls=":", lw=1.3, label=r"3 $\sigma$ contour"
    )
    fermi = mlines.Line2D(
        [], [], color="#21ABCD", marker="", ls="-", lw=2., label="Fermi-LAT"
    )
    magic = mlines.Line2D(
        [], [], color="#FF9933", marker="", ls="-", lw=2., label="MAGIC"
    )
    hess = mlines.Line2D(
        [], [], color="#5A4FCF", marker="", ls="-", lw=2., label="H.E.S.S."
    )
    fact = mlines.Line2D(
        [], [], color="#5CC184", marker="", ls="-", lw=2., label="FACT"
    )
    veritas = mlines.Line2D(
        [], [], color="#702963", marker="", ls="-", lw=2., label="VERITAS"
    )
    joint = mlines.Line2D(
        [], [], color="crimson", marker="", ls="-", lw=2., label="joint fit"
    )
    axarr[2].legend(
        handles=[sigma_1, sigma_2, sigma_3, fermi, magic, hess, fact, veritas, joint],
        loc=3,
        fontsize=12,
    )
    # axarr[2].set_axis_off()

    plt.tight_layout()
    filename = "results/figures/sherpa_logparabola_contour.png"
    fig.savefig(filename)
    log.info(f"Writing {filename}")
    fig.savefig(filename)

def plot_iminuit_contours():
    """plot the confidence contours obtained from the sherpa fit
    """
    log.info("plotting parameters contours obtained from iminuit")
    # where to take the results, configurations for the individual butterflies
    instruments = ["fermi", "magic", "veritas", "fact", "hess", "joint"]
    labels = ["Fermi-LAT", "MAGIC", "VERITAS", "FACT", "H.E.S.S.", "joint fit"]
    colors = COLORS
    lss = ["--", "--", "--", "--", "--", "-"]

    fig, axarr = plt.subplots(1, 3, figsize=(16, 5))

    # with one loop we realize all the contour plots
    for instrument, label, color, ls in zip(instruments, labels, colors, lss):
        path = config.repo_path / f"results/fit/gammapy/{instrument}"

        contours_path = path / "fit_1.0_sigma_contours_logparabola.npy"
        results_path = path / "fit_results_logparabola.yaml"

        if not path.exists():
            log.warning(f"Missing: {path} . Skipping.")
            continue

        # load the contours and the results of the fit
        contours = np.load(contours_path).tolist()
        results = load_yaml(results_path)
        # true values to be plotted
        amplitude = float(results["parameters"][0]["value"])
        alpha = float(results["parameters"][2]["value"])
        beta = float(results["parameters"][3]["value"])

        # amplitude vs alpha
        amplitude_alpha = contours["contour_amplitude_alpha"]
        axarr[0].plot(
            amplitude_alpha["amplitude"] * 10,
            amplitude_alpha["alpha"],
            marker="",
            ls="-",
            lw=2.5,
            color=color,
        )
        # plot actual value
        axarr[0].plot(
            amplitude * 1e11, alpha, marker="X", markersize=7, color=color, lw=2.5
        )
        axarr[0].set_xlabel(
            r"$\phi_0 \,/\,(10^{-11}\,{\rm TeV} \, {\rm cm}^{-2} {\rm s}^{-1})$",
            size=FONTSIZE_CONTOURS,
        )
        axarr[0].set_ylabel(r"$\Gamma$", size=FONTSIZE_CONTOURS)
        # make axis thicker
        for axis in ["top", "bottom", "left", "right"]:
            axarr[0].spines[axis].set_linewidth(2.5)
        axarr[0].set_yticks([2.2, 2.4, 2.6, 2.8])
        axarr[0].set_ylim([2.1, 2.9])
        axarr[0].set_xticks([3, 4, 5])
        axarr[0].set_xlim([2.8, 5.2])
        axarr[0].tick_params(
            "both", length=7, width=1.6, which="major", labelsize=FONTSIZE_CONTOURS
        )
        axarr[0].tick_params(
            "both", length=4, width=1.6, which="minor", labelsize=FONTSIZE_CONTOURS
        )

        # amplitude vs beta
        amplitude_beta = contours["contour_amplitude_beta"]
        axarr[1].plot(
            amplitude_beta["amplitude"] * 10,
            # contour have a scale factor of 1e-10, parameters are in units of 1e-11
            amplitude_beta["beta"],
            marker="",
            ls="-",
            lw=2.5,
            color=color,
        )
        # plot actual value
        axarr[1].plot(amplitude * 1e11, beta, marker="X", markersize=7, color=color)
        axarr[1].set_xlabel(
            r"$\phi_0 \,/\,(10^{-11}\,{\rm TeV} \, {\rm cm}^{-2} {\rm s}^{-1})$",
            size=FONTSIZE_CONTOURS,
        )
        axarr[1].set_ylabel(r"$\beta$", size=FONTSIZE_CONTOURS)
        axarr[1].set_xticks([3, 4, 5])
        axarr[1].set_xlim([2.8, 5.2])
        axarr[1].set_yticks([0.2, 0.4, 0.6])
        axarr[1].set_ylim([0.0, 0.8])
        # make axis thicker
        for axis in ["top", "bottom", "left", "right"]:
            axarr[1].spines[axis].set_linewidth(2.5)
        axarr[1].tick_params(
            "both", length=7, width=1.6, which="major", labelsize=FONTSIZE_CONTOURS
        )
        axarr[1].tick_params(
            "both", length=4, width=1.6, which="minor", labelsize=FONTSIZE_CONTOURS
        )

        # alpha vs beta
        alpha_beta = contours["contour_alpha_beta"]
        axarr[2].plot(
            alpha_beta["alpha"],
            alpha_beta["beta"],
            marker="",
            ls="-",
            lw=2.5,
            color=color,
        )
        # plot actual value
        axarr[2].plot(alpha, beta, marker="X", markersize=7, color=color)
        axarr[2].set_xlabel(r"$\Gamma$", size=FONTSIZE_CONTOURS)
        axarr[2].set_ylabel(r"$\beta$", size=FONTSIZE_CONTOURS)
        axarr[2].set_xticks([2.2, 2.4, 2.6, 2.8])
        axarr[2].set_xlim([2.1, 2.9])
        axarr[2].set_yticks([0.2, 0.4, 0.6])
        axarr[2].set_ylim([0.0, 0.8])
        # make axis thicker
        for axis in ["top", "bottom", "left", "right"]:
            axarr[2].spines[axis].set_linewidth(2.5)
        axarr[2].tick_params(
            "both", length=7, width=1.6, which="major", labelsize=FONTSIZE_CONTOURS
        )
        axarr[2].tick_params(
            "both", length=4, width=1.6, which="minor", labelsize=FONTSIZE_CONTOURS
        )

    # legend
    import matplotlib.lines as mlines

    fermi = mlines.Line2D(
        [], [], color=COLORS[0], marker="", ls="-", lw=2.5, label="Fermi-LAT"
    )
    magic = mlines.Line2D(
        [], [], color=COLORS[1], marker="", ls="-", lw=2.5, label="MAGIC"
    )
    veritas = mlines.Line2D(
        [], [], color=COLORS[2], marker="", ls="-", lw=2.5, label="VERITAS"
    )
    fact = mlines.Line2D(
        [], [], color=COLORS[3], marker="", ls="-", lw=2.5, label="FACT"
    )
    hess = mlines.Line2D(
        [], [], color=COLORS[4], marker="", ls="-", lw=2.5, label="H.E.S.S."
    )
    joint = mlines.Line2D(
        [], [], color=COLORS[5], marker="", ls="-", lw=2.5, label="joint fit"
    )

    box = axarr[2].get_position()
    axarr[2].set_position([box.x0, box.y0, box.width * 0.97, box.height])
    # plot the legend on top of the central plot
    axarr[2].legend(
        handles=[fermi, magic, veritas, fact, hess, joint],
        loc="center left",
        fontsize=FONTSIZE_CONTOURS,
        bbox_to_anchor=(1., 0.5),
    )

    plt.tight_layout()
    filename = "results/figures/iminuit_logparabola_contour.png"
    filename_pdf = "results/figures/iminuit_logparabola_contour.pdf"
    fig.savefig(filename)
    log.info(f"Writing {filename}")
    fig.savefig(filename)
    fig.savefig(filename_pdf)

def run_butterfly_stat(which):
    """run the estimation of the butterfly"""
    if which in ["fermi", "all"]:
        butterfly_stat("fermi")
    if which in ["magic", "all"]:
        butterfly_stat("magic")
    if which in ["veritas", "all"]:
        butterfly_stat("veritas")
    if which in ["fact", "all"]:
        butterfly_stat("fact")
    if which in ["hess", "all"]:
        butterfly_stat("hess")
    if which in ["joint", "all"]:
        butterfly_stat("joint")

def butterfly_stat(which, tool="gammapy"):
    """plot an example figure showing how to correctly estimate a butterfly per each dataset
    We do it by default with the gammapy (iminuit) results
    """
    # the sampling in energy shall be consistent with the nbins of the spectrum extraction
    if which == "fermi":
        num_energy_points = 20
    else:
        num_energy_points = 60
    energies, flux_min, flux_max = stat_errorband(
        which=which, tool=tool, dim_sample=500, energy_points=num_energy_points, sigma=1
    )

    # plot best fit model
    fig, ax = plt.subplots()
    # best fit parameters for this fit
    result_file = f"results/fit/{tool}/{which}/fit_results_logparabola.yaml"
    results = load_yaml(result_file)

    parameters = results["parameters"]
    covariance = results["covariance"]

    # best fit parameters
    amplitude = parameters[0]["value"] * u.Unit(parameters[0]["unit"])
    reference = parameters[1]["value"] * u.Unit(parameters[1]["unit"])
    alpha = parameters[2]["value"] * u.Unit(parameters[2]["unit"])
    beta = parameters[3]["value"] * u.Unit(parameters[3]["unit"])

    # gammapy model with best-fit parameters
    model_lp = Log10Parabola(
        amplitude=amplitude, reference=reference, alpha=alpha, beta=beta
    )

    # set the covariance matrix from the output of the fit, this is needed to propagate the error
    model_lp.parameters.covariance = np.asarray(covariance)

    # plot some model representing the sampling
    # dictionary from the multivariate sampling
    path = Path(
        f"{config.repo_path}/results/debug/stat-err/{tool}/{which}/multivariate_sampling_fluxes.yaml"
    )
    sampled_dict = load_yaml(path)
    sampled_amplitude = sampled_dict["sampled_amplitude"]
    sampled_alpha = sampled_dict["sampled_alpha"]
    sampled_beta = sampled_dict["sampled_beta"]

    for (_ampl, _alpha, _beta) in zip(
        sampled_amplitude[:100], sampled_alpha[:100], sampled_beta[:100]
    ):
        _amplitude = _ampl * u.Unit(parameters[0]["unit"])
        _alpha = _alpha * u.Unit(parameters[2]["unit"])
        _beta = _beta * u.Unit(parameters[3]["unit"])

        _model_lp = Log10Parabola(
            amplitude=_amplitude, reference=reference, alpha=_alpha, beta=_beta
        )
        _model_lp.plot(
            energy_range=[energies[0], energies[-1]],
            energy_unit="TeV",
            flux_unit="cm-2 s-1 erg-1",
            energy_power=2,
            color="gray",
            lw=0.8,
            alpha=0.8,
            ax=ax,
        )

    # plot the 68 % containment correction
    ax.plot(energies, (energies ** 2 * flux_min).to("erg cm-2 s-1"), color="k", lw=2.5)
    ax.plot(
        energies,
        (energies ** 2 * flux_max).to("erg cm-2 s-1"),
        color="k",
        lw=2.5,
        label="68% containment" + "\n" + "multivariate sampling",
    )

    model_lp.plot(
        energy_range=[energies[0], energies[-1]],
        energy_unit="TeV",
        flux_unit="cm-2 s-1 erg-1",
        energy_power=2,
        color="k",
        lw=2.2,
        ls="--",
        label="best fit model",
        ax=ax,
    )

    # plt.title(f'{which} dataset')
    for axis in ["top", "bottom", "left", "right"]:
        ax.spines[axis].set_linewidth(2.)
    ax.tick_params("both", length=7, width=1.6, which="major", labelsize=FONTSIZE)
    ax.tick_params("both", length=4, width=1.6, which="minor", labelsize=FONTSIZE)
    ax.set_ylabel(SED_UNIT_LABEL, size=FONTSIZE)
    ax.set_xlabel(E_UNIT_LABEL, size=FONTSIZE)
    ax.set_xscale("log")
    ax.set_yscale("log")
    # ax.set_ylim([1e-12, 2e-10])
    ax.legend(fontsize=FONTSIZE)
    plt.tight_layout()

    figname = f"{config.repo_path}/results/figures/stat_err/{tool}_{which}_butterfly_comparison.png"
    figname_pdf = f"{config.repo_path}/results/figures/stat_err/{tool}_{which}_butterfly_comparison.pdf"
    logging.info(f"saving {figname}")
    fig.savefig(figname)
    fig.savefig(figname_pdf)

    log.info("making debug plots with histograms of the sampled parameters and fluxes")
    fig2, ax2 = plt.subplots(1, 3, sharey=True)
    nbins_params = int(len(sampled_amplitude) / 20)
    ax2[0].hist(sampled_amplitude, bins=nbins_params, color="lightgray")
    ax2[0].axvline(amplitude.value, lw=1.5, color="k")
    ax2[0].axvline(
        amplitude.value - np.sqrt(covariance[0][0]), lw=1.5, ls="--", color="k"
    )
    ax2[0].axvline(
        amplitude.value + np.sqrt(covariance[0][0]), lw=1.5, ls="--", color="k"
    )
    ax2[0].set_xlabel(
        r"$f_0 (\mathrm{TeV}^{-1} \, \mathrm{cm}^{-2} \, \mathrm{s}^{-1})$", labelpad=9
    )

    ax2[1].hist(sampled_alpha, bins=int(len(sampled_alpha) / 10), color="lightgray")
    ax2[1].axvline(alpha.value, lw=1.5, color="k")
    ax2[1].axvline(alpha.value - np.sqrt(covariance[2][2]), lw=1.5, ls="--", color="k")
    ax2[1].axvline(alpha.value + np.sqrt(covariance[2][2]), lw=1.5, ls="--", color="k")
    ax2[1].set_xlabel(r"$\Gamma$")

    ax2[2].hist(
        sampled_beta,
        bins=int(len(sampled_beta) / 10),
        color="lightgray",
        label="sampled",
    )
    ax2[2].axvline(beta.value, lw=1.5, color="k", label="best fit")
    ax2[2].axvline(
        beta.value - np.sqrt(covariance[3][3]),
        lw=1.5,
        ls="--",
        color="k",
        label=r"best fit $\pm$ 1 $\sigma$",
    )
    ax2[2].axvline(beta.value + np.sqrt(covariance[3][3]), lw=1.5, ls="--", color="k")
    ax2[2].legend()
    ax2[2].set_xlabel(r"$\beta$")

    figname2 = f"{config.repo_path}/results/figures/stat_err/{tool}_{which}_sampled_parameters.png"
    logging.info(f"saving {figname2}")
    fig2.savefig(figname2)

    # sampled fluxes
    fig3, ax3 = plt.subplots(1, 3, sharey=True, figsize=(12, 8))
    for i, ene in enumerate(["emin", "emid", "emax"]):
        flux_dict = sampled_dict["sampled_fluxes"][ene]
        energy = sampled_dict["sampled_fluxes"][ene]["value"] * u.Unit(
            sampled_dict["sampled_fluxes"]["energy_unit"]
        )
        log.info(f"building histogram of sampled fluxes for {energy}")
        nbins_fluxes = int(len(flux_dict["fluxes"]) / 20)
        bins_fluxes = np.logspace(
            np.log10(np.min(flux_dict["fluxes"])),
            np.log10(np.max(flux_dict["fluxes"])),
            nbins_fluxes,
        )
        ax3[i].hist(
            flux_dict["fluxes"],
            bins=bins_fluxes,
            color="lightgray",
            label="sampled fluxes",
        )
        ax3[i].axvline(
            np.mean(flux_dict["fluxes"]),
            lw=1.5,
            ls="-",
            color="k",
            label="mean sampled",
        )
        ax3[i].axvline(
            flux_dict["flux_quantiles"][0],
            lw=1.5,
            ls="--",
            color="k",
            label="68% containment sampled",
        )
        ax3[i].axvline(flux_dict["flux_quantiles"][1], lw=1.5, ls="--", color="k")
        ax3[i].set_xlabel(
            r"$F (\mathrm{TeV}^{-1} \, \mathrm{cm}^{-2} \, \mathrm{s}^{-1})$",
            labelpad=9,
        )
        # evaluate normal model and error
        fit_value = model_lp(energy).to("TeV-1 cm-2 s-1").value
        fit_value_err = model_lp.evaluate_error(energy).to("TeV-1 cm-2 s-1").value
        ax3[i].axvline(
            fit_value, lw=2, ls="-", color="crimson", label="best fit result"
        )
        ax3[i].axvline(
            fit_value + fit_value_err[1],
            lw=1.5,
            ls="--",
            color="crimson",
            label="error propagation",
        )
        ax3[i].axvline(fit_value - fit_value_err[1], lw=1.5, ls="--", color="crimson")
        ax3[i].set_title("E = {:.2f}".format(energy))
        ax3[i].legend()
        ax3[i].set_xscale("log")

    plt.tight_layout()
    figname3 = (
        f"{config.repo_path}/results/figures/stat_err/{tool}_{which}_sampled_fluxes.png"
    )
    logging.info(f"saving {figname3}")
    fig3.savefig(figname3)

def counts_histogram(predicted=False):
    """function to plot the excesses per dataset and compare them with the predicted counts
    if predicted == True will shwo the predicted counts from the results of the fit (for debug purpose)
    """
    log.info("loading the results from the joint fit to predict the counts")
    results = load_yaml(
        f"{config.repo_path}/results/fit/gammapy/joint/fit_results_logparabola.yaml"
    )
    parameters = results["parameters"]

    model_lp = LogParabola.from_log10(
        amplitude=parameters[0]["value"] * u.Unit(parameters[0]["unit"]),
        reference=parameters[1]["value"] * u.Unit(parameters[1]["unit"]),
        alpha=parameters[2]["value"] * u.Unit(parameters[2]["unit"]),
        beta=parameters[3]["value"] * u.Unit(parameters[3]["unit"]),
    )

    # defining the figure
    dict_color = {
        "fermi": COLORS[0],
        "magic": COLORS[1],
        "veritas": COLORS[2],
        "fact": COLORS[3],
        "hess": COLORS[4],
    }
    fig, ax = plt.subplots()

    for which in config.all_datasets:
        log.info(f"predicting counts for {which} dataset")
        dataset = config.get_dataset(which)
        obs = dataset.get_SpectrumObservationList().stack()
        cts_pred = CountsPredictor(
            model=model_lp, aeff=obs.aeff, edisp=obs.edisp, livetime=obs.livetime
        )
        cts_pred.run()

        e_max = dataset.energy_range[1].to("TeV").value
        e_min = dataset.energy_range[0].to("TeV").value

        kwargs_mdl = dict(ls=":", range=(e_min, e_max), lw=2.2, color=dict_color[which])
        kwargs_data = dict(
            ls="-", range=(e_min, e_max), lw=2.2, color=dict_color[which]
        )

        # CountsSpectrum with observed and predicted excesses
        ex_pred = cts_pred.npred
        ex_obs = obs.excess_vector
        # if it is an IACT rebin the counts before plotting
        if which != "fermi":
            ex_pred = ex_pred.rebin(2)
            ex_obs = ex_obs.rebin(2)

        if predicted:  # if you want to display the predicted counts
            ex_pred.plot_hist(ax, **kwargs_mdl)
        ex_obs.plot_hist(ax, **kwargs_data)

    # custom legend
    legend_observed = mlines.Line2D(
        [], [], color="gray", marker="", ls="-", lw=2, label="observed"
    )
    legend_expected = mlines.Line2D(
        [], [], color="gray", marker="", ls=":", lw=2, label="expected"
    )
    legend_fermi = mlines.Line2D(
        [], [], color=COLORS[0], marker="", ls="-", lw=2, label="Fermi-LAT"
    )
    legend_magic = mlines.Line2D(
        [], [], color=COLORS[1], marker="", ls="-", lw=2, label="MAGIC"
    )
    legend_veritas = mlines.Line2D(
        [], [], color=COLORS[2], marker="", ls="-", lw=2, label="VERITAS"
    )
    legend_fact = mlines.Line2D(
        [], [], color=COLORS[3], marker="", ls="-", lw=2, label="FACT"
    )
    legend_hess = mlines.Line2D(
        [], [], color=COLORS[4], marker="", ls="-", lw=2, label="H.E.S.S."
    )
    legend_handles = [
        legend_fermi,
        legend_magic,
        legend_veritas,
        legend_fact,
        legend_hess,
    ]
    if predicted:  # if you want to display the predicted counts
        legend_handles = [legend_observed, legend_expected] + legend_handles

    ax.legend(handles=legend_handles, fontsize=FONTSIZE)

    ax.set_xscale("log")
    ax.set_ylabel("Excess counts", size=FONTSIZE)
    ax.set_xlabel(E_UNIT_LABEL, size=FONTSIZE)

    # make axis thicker
    for axis in ["top", "bottom", "left", "right"]:
        ax.spines[axis].set_linewidth(1.6)
    ax.tick_params("both", length=7, width=1.6, which="major", labelsize=FONTSIZE)
    ax.tick_params("both", length=4, width=1.6, which="minor", labelsize=FONTSIZE)

    plt.tight_layout()

    filename = f"{config.repo_path}/results/figures/counts_spectra.png"
    filename_pdf = f"{config.repo_path}/results/figures/counts_spectra.pdf"
    log.info(f"saving figure in {filename}")
    fig.savefig(filename)
    fig.savefig(filename_pdf)

def butterfly_syst():
    """representation of the joint fit result w/ systematics uncertainty"""
    log.info("plotting syst. + stat. butterfly")
    # reading the joint stat butterfly values
    stat_table_path = Path(
        f"{config.repo_path}/results/figures/stat_err/joint_flux_errorband.dat"
    )
    log.info(f"reading butterfly values from {stat_table_path}")
    stat_table = Table.read(stat_table_path, format="ascii.ecsv")
    stat_energies = stat_table["energies"].data * stat_table["energies"].unit
    stat_flux_min = stat_table["flux_lo"].data * stat_table["flux_lo"].unit
    stat_flux_max = stat_table["flux_hi"].data * stat_table["flux_hi"].unit
    # load result of stat fit
    stat_result_file = f"results/fit/gammapy/joint/fit_results_logparabola.yaml"
    stat_results = load_yaml(stat_result_file)

    stat_parameters = stat_results["parameters"]

    # best fit parameters
    stat_amplitude = stat_parameters[0]["value"] * u.Unit(stat_parameters[0]["unit"])
    stat_reference = stat_parameters[1]["value"] * u.Unit(stat_parameters[1]["unit"])
    stat_alpha = stat_parameters[2]["value"] * u.Unit(stat_parameters[2]["unit"])
    stat_beta = stat_parameters[3]["value"] * u.Unit(stat_parameters[3]["unit"])

    # reading the joint stat butterfly values
    syst_table_path = Path(
        f"{config.repo_path}/results/figures/syst_err/joint_flux_errorband.dat"
    )
    log.info(f"reading butterfly values from {syst_table_path}")
    syst_table = Table.read(syst_table_path, format="ascii.ecsv")
    syst_energies = syst_table["energies"].data * syst_table["energies"].unit
    syst_flux_min = syst_table["flux_lo"].data * syst_table["flux_lo"].unit
    syst_flux_max = syst_table["flux_hi"].data * syst_table["flux_hi"].unit

    # load result of syst fit
    syst_result_file = (
        f"results/fit/gammapy/joint/fit_results_logparabola_energy_scale.yaml"
    )
    syst_results = load_yaml(syst_result_file)

    syst_parameters = syst_results["parameters"]

    # best fit parameters
    syst_amplitude = syst_parameters[0]["value"] * u.Unit(syst_parameters[0]["unit"])
    syst_reference = syst_parameters[1]["value"] * u.Unit(syst_parameters[1]["unit"])
    syst_alpha = syst_parameters[2]["value"] * u.Unit(syst_parameters[2]["unit"])
    syst_beta = syst_parameters[3]["value"] * u.Unit(syst_parameters[3]["unit"])

    fig = plt.figure()
    gs = gridspec.GridSpec(2, 1, height_ratios=[2.5, 1], wspace=0.)
    ax0 = plt.subplot(gs[0])
    ax1 = plt.subplot(gs[1], sharex=ax0)

    ax0.fill_between(
        syst_energies.to("TeV"),
        (syst_energies ** 2 * syst_flux_min).to("erg cm-2 s-1"),
        (syst_energies ** 2 * syst_flux_max).to("erg cm-2 s-1"),
        color="#002E63",
        alpha=0.4,
        label="stat. + syst.",
    )

    ax0.fill_between(
        stat_energies.to("TeV"),
        (stat_energies ** 2 * stat_flux_min).to("erg cm-2 s-1"),
        (stat_energies ** 2 * stat_flux_max).to("erg cm-2 s-1"),
        color="crimson",
        alpha=0.4,
        label="stat. only",
    )

    stat_flux = Log10Parabola.evaluate(
        stat_energies, stat_amplitude, stat_reference, stat_alpha, stat_beta
    )

    ax0.plot(
        stat_energies.to("TeV"),
        (stat_energies ** 2 * stat_flux).to("erg cm-2 s-1"),
        color="crimson",
        lw=2.2,
        ls="-",
    )

    syst_flux = Log10Parabola.evaluate(
        syst_energies, syst_amplitude, syst_reference, syst_alpha, syst_beta
    )

    ax0.plot(
        syst_energies.to("TeV"),
        (syst_energies ** 2 * syst_flux).to("erg cm-2 s-1"),
        color="#002E63",
        lw=2.2,
        ls="-",
    )

    ax0.set_xscale("log")
    ax0.set_yscale("log")
    ax0.set_ylabel(SED_UNIT_LABEL, size=FONTSIZE)
    ax0.legend(fontsize=FONTSIZE, loc=3)
    # make axis thicker
    for axis in ["top", "bottom", "left", "right"]:
        ax0.spines[axis].set_linewidth(1.6)
    ax0.tick_params("both", length=7, width=1.6, which="major", labelsize=FONTSIZE)
    ax0.tick_params("both", length=4, width=1.6, which="minor", labelsize=FONTSIZE)

    # plot the flux ratios
    ax1.fill_between(
        syst_energies.to("TeV"),
        (syst_flux_min - stat_flux) / stat_flux,
        (syst_flux_max - stat_flux) / stat_flux,
        color="#002E63",
        alpha=0.4,
        label="stat. + syst.",
    )

    ax1.fill_between(
        stat_energies.to("TeV"),
        (stat_flux_min - stat_flux) / stat_flux,
        (stat_flux_max - stat_flux) / stat_flux,
        color="crimson",
        alpha=0.4,
        label="stat. only",
    )

    ax1.plot(
        stat_energies.to("TeV"),
        (stat_flux - stat_flux) / stat_flux,
        color="crimson",
        lw=2.2,
        ls="-",
    )

    ax1.plot(
        syst_energies.to("TeV"),
        (syst_flux - stat_flux) / stat_flux,
        color="#002E63",
        lw=2.2,
        ls="-",
    )

    ax1.set_xscale("log")
    ax1.set_ylim([-0.35, 0.35])
    ax1.set_xlabel(E_UNIT_LABEL, size=FONTSIZE)
    ax1.set_ylabel("fractional diff. \n to stat.", size=FONTSIZE)
    # make axis thicker
    for axis in ["top", "bottom", "left", "right"]:
        ax1.spines[axis].set_linewidth(1.6)
    ax1.tick_params("both", length=7, width=1.6, which="major", labelsize=FONTSIZE)
    ax1.tick_params("both", length=4, width=1.6, which="minor", labelsize=FONTSIZE)

    plt.tight_layout()

    filename = "results/figures/crab_sed_joint_fit_syst.png"
    filename_pdf = "results/figures/crab_sed_joint_fit_syst.pdf"
    log.info(f"Writing {filename}")
    fig.savefig(filename)
    fig.savefig(filename_pdf)

def syst_contour():
    """plot a comparison of the contours obtained with the wstat statistics
    and with the stat + syst likelihood"""

    fig, axarr = plt.subplots(1, 3, figsize=(16, 5))

    # first plot in light gray the stat contours
    instruments = ["fermi", "magic", "veritas", "fact", "hess", "joint"]

    # with one loop we realize all the contour plots
    color = "lightgray"
    for instrument in instruments:
        path = f"{config.repo_path}/results/fit/gammapy/{instrument}"

        contours_path = f"{path}/fit_1.0_sigma_contours_logparabola.npy"
        results_path = f"{path}/fit_results_logparabola.yaml"

        # load the contours and the results of the fit
        contours = np.load(contours_path).tolist()
        results = load_yaml(results_path)
        # true values to be plotted
        amplitude = float(results["parameters"][0]["value"])
        alpha = float(results["parameters"][2]["value"])
        beta = float(results["parameters"][3]["value"])

        # amplitude vs alpha
        amplitude_alpha = contours["contour_amplitude_alpha"]
        axarr[0].plot(
            amplitude_alpha["amplitude"] * 10,
            amplitude_alpha["alpha"],
            marker="",
            ls="-",
            lw=2.5,
            color=color,
        )
        # plot actual value
        axarr[0].plot(
            amplitude * 1e11, alpha, marker="X", markersize=7, color=color, lw=2.5
        )

        # amplitude vs beta
        amplitude_beta = contours["contour_amplitude_beta"]
        axarr[1].plot(
            amplitude_beta["amplitude"] * 10,
            # contour have a scale factor of 1e-10, parameters are in units of 1e-11
            amplitude_beta["beta"],
            marker="",
            ls="-",
            lw=2.5,
            color=color,
        )
        # plot actual value
        axarr[1].plot(amplitude * 1e11, beta, marker="X", markersize=7, color=color)

        # alpha vs beta
        alpha_beta = contours["contour_alpha_beta"]
        axarr[2].plot(
            alpha_beta["alpha"],
            alpha_beta["beta"],
            marker="",
            ls="-",
            lw=2.5,
            color=color,
        )
        # plot actual value
        axarr[2].plot(alpha, beta, marker="X", markersize=7, color=color)

    # plotting stat vs syst result on top of instrument-wise contours in gray
    path = config.repo_path / f"results/fit/gammapy/joint"

    stat_contours_path = f"{path}/fit_1.0_sigma_contours_logparabola.npy"
    stat_results_path = f"{path}/fit_results_logparabola.yaml"

    syst_contours_path = f"{path}/fit_1.0_sigma_contours_logparabola_energy_scale.npy"
    syst_results_path = f"{path}/fit_results_logparabola_energy_scale.yaml"

    for (contours_path, results_path, color, label) in zip(
        [stat_contours_path, syst_contours_path],
        [stat_results_path, syst_results_path],
        ["crimson", "#002E63"],
        ["stat. only", "stat. + syst."],
    ):
        # load the contours and the results of the fit
        contours = np.load(contours_path).tolist()
        results = load_yaml(results_path)
        # true values to be plotted
        amplitude = float(results["parameters"][0]["value"])
        alpha = float(results["parameters"][2]["value"])
        beta = float(results["parameters"][3]["value"])

        # amplitude vs alpha
        amplitude_alpha = contours["contour_amplitude_alpha"]
        if label == "stat. only":
            axarr[0].plot(
                amplitude_alpha["amplitude"] * 10,
                # gammapy contour have a scale factor of 1e-10, parameters are in units of 1e-11
                amplitude_alpha["alpha"],
                marker="",
                ls="-",
                lw=2.5,
                color=color,
            )
        else:
            axarr[0].plot(
                amplitude_alpha["amplitude"] * 1e11,
                amplitude_alpha["alpha"],
                marker="",
                ls="-",
                lw=2.5,
                color=color,
            )

        # plot actual value
        axarr[0].plot(
            amplitude * 1e11, alpha, marker="X", markersize=7, color=color, lw=2.5
        )

        # amplitude vs beta
        amplitude_beta = contours["contour_amplitude_beta"]
        if label == "stat. only":
            axarr[1].plot(
                amplitude_beta["amplitude"] * 10,
                # gammapy contour have a scale factor of 1e-10, parameters are in units of 1e-11
                amplitude_beta["beta"],
                marker="",
                ls="-",
                lw=2.5,
                color=color,
            )
        else:
            axarr[1].plot(
                amplitude_beta["amplitude"] * 1e11,
                amplitude_beta["beta"],
                marker="",
                ls="-",
                lw=2.5,
                color=color,
            )

        # plot actual value
        axarr[1].plot(amplitude * 1e11, beta, marker="X", markersize=7, color=color)

        # alpha vs beta
        alpha_beta = contours["contour_alpha_beta"]
        axarr[2].plot(
            alpha_beta["alpha"],
            alpha_beta["beta"],
            marker="",
            ls="-",
            lw=2.5,
            color=color,
        )
        # plot actual value
        axarr[2].plot(alpha, beta, marker="X", markersize=7, color=color)

    # legend
    import matplotlib.lines as mlines

    stat_label = mlines.Line2D(
        [],
        [],
        color="crimson",
        marker="",
        ls="-",
        lw=2.5,
        label="joint fit \n stat. only",
    )
    syst_label = mlines.Line2D(
        [],
        [],
        color="#002E63",
        marker="",
        ls="-",
        lw=2.5,
        label="joint fit \n stat. + syst.",
    )
    single_label = mlines.Line2D(
        [],
        [],
        color="lightgray",
        marker="",
        ls="-",
        lw=2.5,
        label="instr. fit \n stat. only",
    )

    axarr[0].set_xlabel(
        r"$\phi_0 \,/\,(10^{-11}\,{\rm TeV} \, {\rm cm}^{-2} {\rm s}^{-1})$",
        size=FONTSIZE_CONTOURS,
    )
    axarr[0].set_ylabel(r"$\Gamma$", size=FONTSIZE_CONTOURS)
    # make axis thicker
    axarr[0].set_yticks([2.2, 2.4, 2.6, 2.8])
    axarr[0].set_ylim([2.1, 2.9])
    axarr[0].set_xticks([3, 4, 5])
    axarr[0].set_xlim([2.8, 5.2])
    for axis in ["top", "bottom", "left", "right"]:
        axarr[0].spines[axis].set_linewidth(2.5)
    axarr[0].tick_params(
        "both", length=7, width=1.6, which="major", labelsize=FONTSIZE_CONTOURS
    )
    axarr[0].tick_params(
        "both", length=4, width=1.6, which="minor", labelsize=FONTSIZE_CONTOURS
    )

    axarr[1].set_xlabel(
        r"$\phi_0 \,/\,(10^{-11}\,{\rm TeV} \, {\rm cm}^{-2} {\rm s}^{-1})$",
        size=FONTSIZE_CONTOURS,
    )
    axarr[1].set_ylabel(r"$\beta$", size=FONTSIZE_CONTOURS)
    # make axis thicker
    axarr[1].set_xticks([3, 4, 5])
    axarr[1].set_xlim([2.8, 5.2])
    axarr[1].set_yticks([0.2, 0.4, 0.6])
    axarr[1].set_ylim([0.0, 0.8])
    for axis in ["top", "bottom", "left", "right"]:
        axarr[1].spines[axis].set_linewidth(2.5)
    axarr[1].tick_params(
        "both", length=7, width=1.6, which="major", labelsize=FONTSIZE_CONTOURS
    )
    axarr[1].tick_params(
        "both", length=4, width=1.6, which="minor", labelsize=FONTSIZE_CONTOURS
    )

    axarr[2].set_xlabel(r"$\Gamma$", size=FONTSIZE_CONTOURS)
    axarr[2].set_ylabel(r"$\beta$", size=FONTSIZE_CONTOURS)
    axarr[2].set_xticks([2.2, 2.4, 2.6, 2.8])
    axarr[2].set_xlim([2.1, 2.9])
    axarr[2].set_yticks([0.2, 0.4, 0.6])
    axarr[2].set_ylim([0.0, 0.8])
    # make axis thicker
    for axis in ["top", "bottom", "left", "right"]:
        axarr[2].spines[axis].set_linewidth(2.5)
    axarr[2].tick_params(
        "both", length=7, width=1.6, which="major", labelsize=FONTSIZE_CONTOURS
    )
    axarr[2].tick_params(
        "both", length=4, width=1.6, which="minor", labelsize=FONTSIZE_CONTOURS
    )

    box = axarr[2].get_position()
    axarr[2].set_position([box.x0, box.y0, box.width * 0.97, box.height])
    # plot the legend on top of the central plot
    axarr[2].legend(
        handles=[stat_label, syst_label, single_label],
        loc="center left",
        fontsize=FONTSIZE_CONTOURS,
        bbox_to_anchor=(1., 0.5),
    )

    plt.tight_layout()
    filename = "results/figures/iminuit_logparabola_energy_scale_contour.png"
    filename_pdf = "results/figures/iminuit_logparabola_energy_scale_contour.pdf"
    fig.savefig(filename)
    log.info(f"Writing {filename}")
    fig.savefig(filename)
    fig.savefig(filename_pdf)


import torch.nn as nn
import torch.nn.functional as F
import torch

import numpy as np

def model_box_to_xy_ori(box,ori_im_shape,prior_w = 0.5, prior_cx=0.5, prior_cy=0.5,size_variance = 0.1,center_variance = 0.1):

    box[:,0] = (box[:,0] * center_variance +  prior_cx * 2.0)/2.0
    box[:,1] = (box[:,1] * center_variance +  prior_cy * 2.0)/2.0
    box[:,2] = torch.exp(box[:,2] * size_variance) * prior_w
    box[:,3] = torch.exp(box[:,3] * size_variance) * prior_w

    box_numpy = box.cpu().numpy()

    box_xyxy = np.zeros((box_numpy.shape[0],4),dtype = 'float32')

    box_xyxy[:,0] = box_numpy[:,0] - box_numpy[:,2]/2.0
    box_xyxy[:,1] = box_numpy[:,1] - box_numpy[:,3]/2.0
    box_xyxy[:,2] = box_numpy[:,0] + box_numpy[:,2]/2.0
    box_xyxy[:,3] = box_numpy[:,1] + box_numpy[:,3]/2.0

    return box_xyxy

def model_box_to_xy_ori_l1(box):

    box_numpy = box.cpu().numpy()

    box_xyxy = np.zeros((box_numpy.shape[0],4),dtype = 'float32')

    box_xyxy[:,0] = box_numpy[:,0] - box_numpy[:,2]/2.0
    box_xyxy[:,1] = box_numpy[:,1] - box_numpy[:,3]/2.0
    box_xyxy[:,2] = box_numpy[:,0] + box_numpy[:,2]/2.0
    box_xyxy[:,3] = box_numpy[:,1] + box_numpy[:,3]/2.0

    return box_xyxy

def norm_box_to_abs(box_norm,ori_im_shape):
    box = box_norm.copy()
    box[0] *=ori_im_shape[1]
    box[1] *=ori_im_shape[0]
    box[2] *=ori_im_shape[1]
    box[3] *=ori_im_shape[0]
    return box

def area_of(left_top, right_bottom) -> torch.Tensor:
    """Compute the areas of rectangles given two corners.

    Args:
        left_top (N, 2): left top corner.
        right_bottom (N, 2): right bottom corner.

    Returns:
        area (N): return the area.
    """
    hw = torch.clamp(right_bottom - left_top, min=0.0)
    return hw[..., 0] * hw[..., 1]

def iou_of(boxes0, boxes1, eps=1e-5):
    """Return intersection-over-union (Jaccard index) of boxes.

    Args:
        boxes0 (N, 4): ground truth boxes.
        boxes1 (N or 1, 4): predicted boxes.
        eps: a small number to avoid 0 as denominator.
    Returns:
        iou (N): IoU values.
    """
    overlap_left_top = torch.max(boxes0[..., :2], boxes1[..., :2])
    overlap_right_bottom = torch.min(boxes0[..., 2:], boxes1[..., 2:])

    overlap_area = area_of(overlap_left_top, overlap_right_bottom)
    area0 = area_of(boxes0[..., :2], boxes0[..., 2:])
    area1 = area_of(boxes1[..., :2], boxes1[..., 2:])
    return overlap_area / (area0 + area1 - overlap_area + eps)

def iou_ori(box_pred,box_gt,ori_im_shape):
    # one box input
    b_pred = torch.from_numpy(box_pred)
    b_gt = torch.from_numpy(box_gt)

    b_gt[:,0] *=ori_im_shape[1]
    b_gt[:,1] *=ori_im_shape[0]
    b_gt[:,2] *=ori_im_shape[1]
    b_gt[:,3] *=ori_im_shape[0]

    b_pred[:,0] *=ori_im_shape[1]
    b_pred[:,1] *=ori_im_shape[0]
    b_pred[:,2] *=ori_im_shape[1]
    b_pred[:,3] *=ori_im_shape[0]

    return iou_of(b_pred,b_gt)

#class Location_Box_Loss(nn.Module):
#    def __init__(self,prior_w = 0.5, prior_cx=0.5, prior_cy=0.5,size_variance = 0.1,center_variance = 0.1):
#        """Implement Loss of box locolization
#
#        smooth L1 regression loss.
#        """
#        super(Location_Box_Loss, self).__init__()
#        self.prior_w = prior_w
#
#        self.prior_cx = prior_cx
#        self.prior_cy = prior_cy
#
#        self.size_variance = size_variance
#        self.center_variance = center_variance
#
#    def forward(self, gt_box, predicted_vec):
#        """
#        # predicted_vec  x_center(-1,1)  /center_variance
#                         y_center(-1,1) /center_variance
#                         w ( -inf, inf)  0-1x, 1-2.7x -1 -0.36x  /size_variance
#        """
#        ww_gt = gt_box[:,2] - gt_box[:,0]
#        ww_gt_tr   = torch.log(ww_gt/self.prior_w)/self.size_variance
#
#        hh_gt = gt_box[:,3] - gt_box[:,1]
#        hh_gt_tr   = torch.log(hh_gt/self.prior_w)/self.size_variance
#
#        xx_gt_tr = ((gt_box[:,2] + gt_box[:,0]) - self.prior_cx * 2.0)/ self.center_variance   # (0.1) -> (-1,1)
#        yy_gt_tr = ((gt_box[:,3] + gt_box[:,1]) - self.prior_cy * 2.0)/ self.center_variance   # (0.1) -> (-1,1)
#
#        loss_xx = F.smooth_l1_loss(predicted_vec[:,0], xx_gt_tr, size_average=True)
#        loss_yy = F.smooth_l1_loss(predicted_vec[:,1], yy_gt_tr, size_average=True)
#
#        loss_ww = F.smooth_l1_loss(predicted_vec[:,2], ww_gt_tr, size_average=True)
#        loss_hh = F.smooth_l1_loss(predicted_vec[:,3], hh_gt_tr, size_average=True)
#
#
#        loss = loss_xx + loss_yy +  loss_ww + loss_hh
#        return loss,loss_xx,loss_yy,loss_ww,loss_hh

class Location_Box_Loss(nn.Module):
    def __init__(self,prior_w = 0.5, prior_cx=0.5, prior_cy=0.5,size_variance = 0.1,center_variance = 0.1):
        """Implement Loss of box locolization

        smooth L1 regression loss.
        """
        super(Location_Box_Loss, self).__init__()
        self.prior_w = prior_w

        self.prior_cx = prior_cx
        self.prior_cy = prior_cy

        self.size_variance = size_variance
        self.center_variance = center_variance

    def forward(self, gt_box, predicted_vec):
        """
        # predicted_vec  x_center(-1,1)  /center_variance
                         y_center(-1,1) /center_variance
                         w ( -inf, inf)  0-1x, 1-2.7x -1 -0.36x  /size_variance
        """
        ww_gt = gt_box[:,2] - gt_box[:,0] # (0,1)
        #ww_gt_tr   = torch.log(ww_gt/self.prior_w)/self.size_variance

        hh_gt = gt_box[:,3] - gt_box[:,1]
        #hh_gt_tr   = torch.log(hh_gt/self.prior_w)/self.size_variance

        xx_gt_tr = (gt_box[:,2] + gt_box[:,0])/2.0  #(0.1)
        yy_gt_tr = (gt_box[:,3] + gt_box[:,1])/2.0   # (0.1)

        #loss_xx = F.smooth_l1_loss(predicted_vec[:,0], xx_gt_tr, size_average=True)
        #loss_yy = F.smooth_l1_loss(predicted_vec[:,1], yy_gt_tr, size_average=True)

        #loss_ww = F.smooth_l1_loss(predicted_vec[:,2], ww_gt, size_average=True)
        #loss_hh = F.smooth_l1_loss(predicted_vec[:,3], hh_gt, size_average=True)

        loss_xx = F.l1_loss(predicted_vec[:,0], xx_gt_tr, size_average=True)
        loss_yy = F.l1_loss(predicted_vec[:,1], yy_gt_tr, size_average=True)

        loss_ww = F.l1_loss(predicted_vec[:,2], ww_gt, size_average=True)
        loss_hh = F.l1_loss(predicted_vec[:,3], hh_gt, size_average=True)

        loss = loss_xx + loss_yy +  loss_ww + loss_hh
        return loss,loss_xx,loss_yy,loss_ww,loss_hh

class Location_Box_CenterNet(nn.Module):
    def __init__(self):
        """Implement Loss of box locolization

        hm focal-loss xywh L1 regression loss.
        """
        super(Location_Box_CenterNet, self).__init__()

    def forward(self, gt_box, pred):

       #score
       hm = gt_box['hm']
       pred_s = pred[:,0:1,...]
       pos_inds = hm.eq(1).float()
       neg_inds = hm.lt(1).float()

       neg_weights = torch.pow(1.0 - hm, 4)
       pos_loss = torch.log(pred_s) * torch.pow(1 - pred_s, 2) * pos_inds
       neg_loss = torch.log(1 - pred_s) * torch.pow(pred_s, 2) * neg_weights * neg_inds

       num_pos  = pos_inds.float().sum()
       pos_loss = pos_loss.sum()
       neg_loss = neg_loss.sum()

       loss_s = -(pos_loss + neg_loss) / num_pos

       dense_xy = gt_box['dense_xy']
       dense_wh = gt_box['dense_wh']
       dense_mask = gt_box['dense_mask']
       pred_xy = pred[:,1:3,...]
       pred_wh = pred[:,3:,...]

       loss_xy = F.l1_loss(pred_xy*dense_mask,dense_xy*dense_mask,size_average=False) / num_pos
       loss_wh = F.l1_loss(pred_wh*dense_mask,dense_wh*dense_mask,size_average=False) / num_pos

       loss = loss_s + loss_xy + loss_wh
       return loss,loss_s,loss_xy,loss_wh