#!/usr/bin/python
# -*- coding: utf-8 -*-
#

import os
import re
import gzip
import sqlite3 as sql
import subprocess
from StringIO import StringIO
from math import degrees, atan2, pi, ceil # , sin, cos
from array import array
from glob import glob

import numpy

from prosci.util.pdb3d import dihedral_angle, dist, sqrdist, vectors2rotation_matrix, axis_angle2rotation_matrix
from prosci.util.residue import Residue, ResidueList
from prosci.util.residue import is_residue_consecutive as is_consecutive
from prosci.util.pdb import Atom, residueCode
from prosci.util.ali import Ali
from prosci.util.geohash import GeometricHash
from prosci.common import join

# Used implicitly in this file:
#
#from prosci.util.residue import Residue

# Database anchor hash granularity
#ANCHOR_GRANULARITY = 1.0

ANCHOR_LENGTH = 2

#STRUC_HASH_GRANULARITY = 0.3

class Data:
  pass

def get_max_loop_length(loop_length):
  return 3.8 * (1+loop_length)

# def is_consecutive(a, b, minlen=1.0, maxlen=2.0):
# #def is_consecutive(residues, a, b):
#     if a.C is None or b.N is None:
#       #raise RuntimeError("Cannot reliably test if two residues are consecutive, purely based on CA distance! Need N and C atoms!")
#       return False
#
#     #d=dist(r.CA, q.CA)
#     #if (d < 3.6 or d > 4.0):
#     d = dist(a.C, b.N)
#     if (d < minlen or d > maxlen):
#       return False
#     return True

def get_loop_structure(dbdir, strucname, start, length):

  def parse(f):
    residues=[]
    r=[]
    i=-1
    for line in f:
      if line.startswith('>'):
        i+=1
        if i >= start+length:
          residues.append(Residue(r))
          break
        if r:
          residues.append(Residue(r))
          r=[]
      elif i>= start:
        r.append(Atom(line))
    return residues

  source = os.path.join(dbdir, "structures", strucname[1:3], strucname)
  f = None
  try:
    #p = subprocess.Popen(["zcat", source + ".atm.gz"], stdout = subprocess.PIPE)
    #f = StringIO(p.communicate()[0])
    #if p.returncode != 0:
    #  raise IOError()
    f = gzip.open(source + ".atm.gz")
    residues = parse(f)
  except IOError:
    f = open(source + ".atm")
    residues = parse(f)
  finally:
    if f is not None:
      f.close()

  return ResidueList(residues)

# # TODO: Doesn't work yet. Needs to get anchor residues too!
# def get_loop_structure_from_decoyname(dbdir, name):
#   "Parse the name of a loop decoy and get its structure from the database"
#
#   strucname, startres, length = name.split("_")
#   startinscode = ""
#   if startres[-1:].isalpha():
#     startinscode = startres[-1:]
#     startres = startres[:-1]
#   startres = int(startres)
#   length = int(length) + 2*ANCHOR_LENGTH
#
#   source = os.path.join(dbdir, "structures", strucname[1:3], strucname)
#   residues=[]
#   r=[]
#   i=-1
#   f = gzip.open(source + ".atm.gz")
#   skip = False
#   try:
#     for line in f:
#       if line.startswith('>'):
#         if i < 0:
#           skip = False
#           continue
#         i+=1
#         if i >= length:
#           residues.append(Residue(r))
#           break
#         if r:
#           residues.append(Residue(r))
#           r=[]
#       elif i < 0:
#         if not skip:
#           a = Atom(line)
#           if a.ires < startres or (a.ires == startres and  a.inscode < startinscode):
#             skip = True
#           else:
#             r.append(Atom(line))
#             i = 0
#       else:
#         r.append(Atom(line))
#     # No need to check if r has not yet been added, as loops never
#     # reach the end of the protein.
#   finally:
#     f.close()
#   return ResidueList(residues)

def get_native_contacts(dbdir, strucname, start, length):
  source = os.path.join(dbdir, "structures", strucname[1:3], strucname)
  a = Ali(source + ".tem")
  dihedrals = a[0]["FREAD dihedral class"].seq[start:start+length]
  contacts = a[0]["FREAD contact class"].seq[start:start+length]
  assert "?" not in dihedrals
  return contacts

def add_oxygens(residues, start=0, end=None, dO=1.23, force=False):
  if end is None:
    end = len(residues) - 1
  for i in xrange(start, end):
    # add missing mainchain oxygens, making them planar
    # with CA, C and the next residue's N
    #
    r = residues[i]
    if r is None:
      continue

    if r.O is None:
      q = residues[i+1]
      if q is None or q.N is None or r.C is None:
        continue

      if not force and not is_consecutive(r, q):
        continue

      O = r.C.copy()
      O.atom = "O"
      if r.C.element:
        O.element = "O"
      O.iatom = r.C.iatom+1

      vO_1 = r.C.xyz - r.CA.xyz
      vO_2 = r.C.xyz - q.N.xyz
      vO_3 = (vO_1 / numpy.linalg.norm(vO_1)) + (vO_2 / numpy.linalg.norm(vO_2))
      O.xyz = r.C.xyz + (vO_3 / numpy.linalg.norm(vO_3)) * dO

      r.O = O

def is_structure_ok(residues, start, end, max_b=float("inf")):
  for i in xrange(start, end):
    r1=residues[i]
    if r1.N is None or r1.CA is None or r1.C is None or r1.O is None or \
       r1.N.b > max_b or r1.CA.b > max_b or r1.C.b > max_b or r1.O.b > max_b:
      return False

  r1=residues[start]
  for i in xrange(start+1, end):
    r2=residues[i]
    if (r2.CA.ires != r1.CA.ires + 1):
      return False
    if not is_consecutive(r1, r2):
      return False
    # d=dist(r1.CA, r2.CA)
    # if (d < 2.8 or d > 4.0):
    #   return False
    r1=r2

  return True

# dihed_phi   = zip([-180, -110,-110,-180,-180,-180,  20,   0,  0],
#                   [   0,    0,   0,-110,-110,   0, 140, 180,180])
# dihed_psi     = zip([ -90,  100,-180, 100,-180,  45, -40,-180, 80],
#                   [  45,  180, -90, 180, -90, 100,  80, -40,180])
# dihed_class =     [   0,    1,   1,   2,   2,   3,   4,   5,  5]

# dihed_phi   = zip([-180, -110,-110,-180,-180,-180,  20,   0],
#                   [   0,    0,   0,-110,-110,   0, 140, 180])
# dihed_psi   = zip([ -90,  100,-180, 100,-180,  80, -40,-180],
#                   [  80,  180, -90, 180, -90, 100,  80, 180])
# dihed_class =     [   0,    1,   1,   2,   2,   3,   4,   5]
# #                     A     B    B    C    C    D    E    F

dihed_phi   = zip([-180, -110,-110,-180,-180,-180,  20,   0],
                  [   0,    0,   0,-110,-110,   0, 140, 180])
dihed_psi   = zip([ -90,  100,-180, 100,-180,  45, -40,-180],
                  [  45,  180, -90, 180, -90, 100,  80, 180])
dihed_class =     [   0,    1,   1,   2,   2,   3,   4,   5]
#                     A     B    B    C    C    D    E    F

def get_dihedral_angle(residues, i):
  r1 = residues[i-1]
  r2 = residues[i]
  r3 = residues[i+1]
  phi = dihedral_angle(r2.N.xyz  - r1.C.xyz, r2.CA.xyz - r2.N.xyz, r2.C.xyz - r2.CA.xyz)
  psi = dihedral_angle(r2.CA.xyz - r2.N.xyz, r2.C.xyz - r2.CA.xyz, r3.N.xyz - r2.C.xyz)
  return (phi, psi)

def angle2class(phi, psi):
  for i in xrange(9):
    minphi,maxphi = dihed_phi[i]
    minpsi,maxpsi = dihed_psi[i]
    if phi <= maxphi and phi >= minphi and psi <= maxpsi and psi >= minpsi:
      return dihed_class[i]
  raise IllegalStateError("illegal dihedral angles : %f, %f" % (phi, psi))

def get_dihedral_class(residues, i):
  phi, psi = get_dihedral_angle(residues, i)
  return angle2class(phi, psi)

def describe_anchors(anchor_n, anchor_c, loop_length):
  max_actual_length = get_max_loop_length(loop_length)

  n1, n2 = anchor_n[-2:]
  c1, c2 = anchor_c[:2]

  # get CA coordinates of anchors
  coords = numpy.array([n1.CA.xyz, n2.CA.xyz, c1.CA.xyz, c2.CA.xyz, n2.O.xyz, c1.O.xyz])

  origin = numpy.copy(coords[1])

  # translate second CA of N-anchor onto origin
  for i,c in enumerate(coords):
    coords[i] -= origin

  #print coords

  # rotate first CA of N-anchor onto Y axis in the negative direction, such that
  # the actual loop goes off in the positive direction along Y axis
  #
#   if numpy.round(coords[0][0], 4) != 0 or numpy.round(coords[0][2], 4) != 0:
#     rotmat1 = vectors2rotation_matrix(coords[0], numpy.array([0.0,-1.0,0.0]))
#     for i,c in enumerate(coords):
#       coords[i] = numpy.dot(c, rotmat1)
#   else:
#     rotmat1 = None

#   print "Experimental rotation matrix:"
#   print vectors2rotation_matrix(numpy.array([-1.0,0.0,0.0]), numpy.array([1.0,0.0,0.0]))
#   print "Experimental rotation matrix2:"
#   exprotmat = axis_angle2rotation_matrix(numpy.array([0.0,1.0,0.0]), pi)
#   print numpy.dot(numpy.array([-1.0,0.0,0.0]), exprotmat)
#   raise RuntimeError()

  # rotate first CA of C-anchor onto X axis in the positive direction
  #
  if numpy.round(coords[2][1], 4) != 0 or numpy.round(coords[2][2], 4) != 0:
    rotmat1 = vectors2rotation_matrix(coords[2], numpy.array([1.0,0.0,0.0]))
    for i,c in enumerate(coords):
      coords[i] = numpy.dot(c, rotmat1)
  elif coords[2][0] < 0:
    rotmat1 = axis_angle2rotation_matrix((0.0,1.0,0.0), pi)
    for i,c in enumerate(coords):
      coords[i] = numpy.dot(c, rotmat1)
  else:
    rotmat1 = None

  #print coords

  # rotate first CA of N-anchor onto X-Y plane, in the negative Y direction, so that the loop always goes off into the positive Y direction.
  if numpy.round(coords[0][2], 4) != 0:
    rotmat2 = vectors2rotation_matrix(numpy.array([0.0, coords[0][1], coords[0][2]]), numpy.array([0.0,-1.0,0.0]))
    for i,c in enumerate(coords):
      coords[i] = numpy.dot(c, rotmat2)
  elif coords[0][1] > 0:
    rotmat2 = axis_angle2rotation_matrix((1.0,0.0,0.0), pi)
    for i,c in enumerate(coords):
      coords[i] = numpy.dot(c, rotmat2)
  else:
    rotmat2 = None

  #print coords

  coords = coords.round(4)

  try:
    assert (coords[1] == (0,0,0)).all()  # second N atom should be the origin
    assert coords[2][1] == 0   # first C atom should be on the X axis
    assert coords[2][2] == 0   # first C atom should be on the X axis
    assert coords[2][0] >  0   # first C atom should be be in the positive X range
    assert coords[0][2] == 0   # first N atom should be on the X-Y plane
    assert coords[0][1] <= 0   # first N atom should be on in the negative Y range
    assert not numpy.isnan(coords).any()   # must not have NaN in coords
  except:
    print "Anchor coordinates:"
    for c in coords:
      print c
      print
    raise

  d = Data()

  # save coordinates
  d.N1 = coords[0]
  # N2 = [0, 0, 0]   # N2 is always on the origin
  d.C1 = coords[2]
  d.C2 = coords[3]

  # compute loop stretch (can just use X coords of first CA of C-anchor, as it's on the X axis)
  stretch = min(1.0, coords[2][0] / max_actual_length)

  # get angle between N anchor and Y axis
  Nanchor_angle = atan2(-coords[0][0], -coords[0][1])

  # get rotation angle of N anchor's O atom around the N anchor's own axis
  #Nanchor_rotmat = numpy.array([[cos(Nanchor_angle), -sin(Nanchor_angle), 0], [sin(Nanchor_angle), cos(Nanchor_angle), 0], [0, 0, 1]]); # matrix rotates N anchor onto Y axis
  Nanchor_rotmat = axis_angle2rotation_matrix((0.0,0.0,1.0), Nanchor_angle)
  Nanchor_O = numpy.dot(coords[4], Nanchor_rotmat)
  Nanchor_twist = atan2(Nanchor_O[0], Nanchor_O[2])

  # get X-Y-Z unit vector from second C atom to first C atom
  Canchor_tilt = coords[2] - coords[3]
  Canchor_tilt /= numpy.linalg.norm(Canchor_tilt) # make unit vector
  Canchor_tilt = numpy.dot(Canchor_tilt, Nanchor_rotmat)

  # get rotation angle of C anchor's O atom around the C anchor's own axis
  Canchor_rotmat = vectors2rotation_matrix(coords[3] - coords[2], numpy.array([0.0,-1.0,0.0]))
  Canchor_O = numpy.dot(coords[5] - coords[2], Canchor_rotmat)
  Canchor_twist = atan2(Canchor_O[0], Canchor_O[2])

  d.stretch = stretch
  d.Nangle = degrees(Nanchor_angle)%180
  d.Ctilt = Canchor_tilt
  d.Ntwist = degrees(Nanchor_twist)%360
  d.Ctwist = degrees(Canchor_twist)%360

  return d, (origin, rotmat1, rotmat2)

# def hash_loop_anchor_numeric(anchor_desc):
#   return (
#     int(anchor_desc.N1[0] / ANCHOR_GRANULARITY),
#     int(anchor_desc.N1[1] / ANCHOR_GRANULARITY),
#     int(anchor_desc.C1[0] / ANCHOR_GRANULARITY),
#     int(anchor_desc.C2[0] / ANCHOR_GRANULARITY),
#     int(anchor_desc.C2[1] / ANCHOR_GRANULARITY),
#     int(anchor_desc.C2[2] / ANCHOR_GRANULARITY)
#     )

# def hash_loop_anchor_string(anchor_desc):
#   return hash_numeric_to_string(hash_loop_anchor_numeric(anchor_desc))

# def hash_numeric_to_string(numeric):
#   return join(".", numeric)

# def hash_string_to_numeric(string):
#   numeric=[]
#   x = string.split('.')
#   for y in x:
#     numeric.append(int(y))
#   return tuple(numeric)

# def hash_structure(coords, transform, bin_size=STRUC_HASH_GRANULARITY):
#   hash=array('c')
#   for i in xrange(len(coords)):
#     c = transform_xyz(coords[i], transform)
#     hash.extend(str(int((c[0] + bin_size/2) / bin_size)))
#     hash.append('.')
#     hash.extend(str(int((c[1] + bin_size/2) / bin_size)))
#     hash.append('.')
#     hash.extend(str(int((c[2] + bin_size/2) / bin_size)))
#     hash.append('.')
#   hash.pop()
#   return hash.tostring()

def transform_xyz(xyz, transform):
  c = xyz - transform[0]
  if transform[1] is not None:
    c = numpy.dot(c, transform[1])
  if transform[2] is not None:
    c = numpy.dot(c, transform[2])
  return c

###########################################################

# def score_decoy(seq, esss_vectors, ascii2index):
#   # Calculate ESSS
#   score=0
#   for i,x in enumerate(seq):
#     score += esss_vectors[i][ascii2index[ord(x)]]
#   return score

def get_db_path(dbdir, loop_length):
  return os.path.join(dbdir, "length%d.sqlite"%loop_length)

def get_db_file_dict(dbdir):
  prefix = os.path.join(dbdir, "length")
  suffix = ".sqlite"
  d={}
  for fname in glob(prefix+"*"+suffix):
    L = int(fname[len(prefix):-len(suffix)])
    d[L] = fname
  return d

def is_loop_length_in_database(dbdir, loop_length):
  return os.path.exists(get_db_path(dbdir, loop_length))

def get_min_loop_length_in_database(dbdir):
  dbfiles = glob(os.path.join(dbdir, "length*.sqlite"))
  if not dbfiles:
    raise RuntimeError("No database found in directory: "+dbdir)
  minlen=1000000
  for fname in dbfiles:
    length = os.path.splitext(os.path.basename(fname))[0].replace("length", "")
    minlen = min(minlen, int(length))
  return minlen

def iterate_database(dbdir, loop_length, esst, anchor, sequence, strict_rmsd_cutoff, score_cutoff):
  q = anchor
  dbfile = get_db_path(dbdir, loop_length)

  # If we cannot possibly meet the score_cutoff with a perfect match to this sequence, reduce cut-off to allow perfect matches to be found
  score_cutoff = min(score_cutoff, esst.get_perfect_score(sequence))

  #print "Score cut-off:", score_cutoff

  tables = esst.tables
  ascii2index = esst.ascii2index
  seqmap = tuple([ascii2index[ord(s)] for s in sequence])

  assert os.path.exists(dbfile), "Database file missing: "+dbfile
  conn = sql.connect(dbfile)

  prevdihed=["", [None]*len(seqmap)]
  def score_sequence(loopseq, dihed):
      if prevdihed[0] != dihed:
        prevdihed[0] = dihed
        for i,x in enumerate(dihed):
          prevdihed[1][i] = tables[int(x)][seqmap[i]]
      score=0
      for i,x in enumerate(loopseq):
        score += prevdihed[1][i][ascii2index[ord(x)]]  # Speed-optimised version
        #score += tables[int(dihed[i])][seqmap[i]][ascii2index[ord(x)]]  # Naive version
      return score

  # DEBUGGING: See if this function throws an exception
  score_sequence(sequence, "0"*len(sequence))

  conn.create_function("score_seq", 2, score_sequence)

  casep1 = q.C1[0]
  casep2 = numpy.linalg.norm(q.C2)
  casep3 = numpy.linalg.norm(q.N1 - q.C1)
  casep4 = numpy.linalg.norm(q.N1 - q.C2)
  min_casep1 = max(0, casep1 - strict_rmsd_cutoff*3)
  max_casep1 =        casep1 + strict_rmsd_cutoff*3
  min_casep2 = max(0, casep2 - strict_rmsd_cutoff*3)
  max_casep2 =        casep2 + strict_rmsd_cutoff*3
  min_casep3 = max(0, casep3 - strict_rmsd_cutoff*3)
  max_casep3 =        casep3 + strict_rmsd_cutoff*3
  min_casep4 = max(0, casep4 - strict_rmsd_cutoff*3)
  max_casep4 =        casep4 + strict_rmsd_cutoff*3

  try:
#      contacts,
    for row in conn.execute("""
    SELECT
      score_seq(sequence, dihedral) AS score,
      dihedral,
      sequence,
      pdbcode,
      start,
      ((casep1-?)*(casep1-?) + (casep2-?)*(casep2-?) + (casep3-?)*(casep3-?) + (casep4-?)*(casep4-?)) / 16 AS internal_rmsd_sq
    FROM
      loops
    WHERE
      casep1 BETWEEN ? AND ?
      AND casep2 BETWEEN ? AND ?
      AND casep3 BETWEEN ? AND ?
      AND casep4 BETWEEN ? AND ?
      AND internal_rmsd_sq <= ?
      AND score >= ?
    """, (casep1, casep1, casep2, casep2, casep3, casep3, casep4, casep4, min_casep1, max_casep1, min_casep2, max_casep2, min_casep3, max_casep3, min_casep4, max_casep4, strict_rmsd_cutoff**2, score_cutoff)):
      d = Data()
      d.score, d.dihedrals, d.seq, d.struc, d.start = row[:5]
      #d.score, d.dihedrals, d.seq, d.struc, d.start, d.native_contacts = row[:6]
      d.internal_rmsd = numpy.sqrt(row[-1])
#       d.casep1_diff = abs(row[5] - casep1)
#       d.casep2_diff = abs(row[6] - casep2)
#       d.casep3_diff = abs(row[7] - casep3)
#       d.casep4_diff = abs(row[8] - casep4)
      yield d
  finally:
    conn.close()

###########################################################

def relabel_loop(decoy_residues, sequence, startnum=0, endnum=0, endinscode="", chain="A", prevatom=None, nextatom=None):
  assert len(decoy_residues) == len(sequence)

  if prevatom:
    chain = prevatom.chain
    startnum = prevatom.ires + 1

  if nextatom:
    if prevatom:
      assert prevatom.chain == nextatom.chain
    else:
      chain = nextatom.chain

    if nextatom.inscode > "A":
      endinscode = chr(ord(nextatom.inscode) - 1)

    if nextatom.inscode:
      endnum = nextatom.ires
    else:
      endnum = nextatom.ires - 1

  inscode = "A"
  for i,r in enumerate(decoy_residues):
    r.set_type(residueCode(sequence[i]))
    if startnum:
      if startnum + i <= endnum:
        for a in r:
          a.ires = startnum + i
          a.inscode = ""
          a.chain = chain
      else:
        for a in r:
          a.ires = endnum
          a.inscode = inscode
          a.chain = chain
        inscode = chr(ord(inscode)+1)

vdw_radii = {'C':1.70, 'N':1.55,'O':1.52,'S':1.80,'P':1.80,'H':1.20}

def is_clash(gh, atoms, decoy, vdw_factor=0.7):
  if vdw_factor <= 0:
    return False

  for r in decoy:
    for a in r:
      if a.atom not in ("N", "CA", "C", "O", "CB"):
        continue
      nbr = gh.get_neighbours(a.xyz, 1, max_dist=3.6*vdw_factor)
      if nbr:
        ix, dst = nbr[0]
        try:
          r1 = vdw_radii[a.atom[0]]
        except KeyError:
          r1 = 1.7
        try:
          r2 = vdw_radii[atoms[ix].atom[0]]
        except KeyError:
          r2 = 1.7
        if dst < (r1 + r2) * vdw_factor:
          return True
  return False

def calculate_rmsd(a, b):
    assert len(a) == len(b)
    total = 0.0
    count = 0
    for i, ra in enumerate(a):
      rb = b[i]
      if ra.N and rb.N:
        total += numpy.linalg.norm(ra.N.xyz - rb.N.xyz)**2
        count += 1
      if ra.CA and rb.CA:
        total += numpy.linalg.norm(ra.CA.xyz - rb.CA.xyz)**2
        count += 1
      if ra.C and rb.C:
        total += numpy.linalg.norm(ra.C.xyz - rb.C.xyz)**2
        count += 1
      if ra.O and rb.O:
        total += numpy.linalg.norm(ra.O.xyz - rb.O.xyz)**2
        count += 1
    return numpy.sqrt(total / count)

def calculate_rmsdCA(a, b):
    assert len(a) == len(b)
    total = 0.0
    for i, ra in enumerate(a):
      rb = b[i]
      total += numpy.linalg.norm(ra.CA.xyz - rb.CA.xyz)**2
    return numpy.sqrt(total / len(a))

def norm_angle(a):
  if a > 180:
    a = 360 - a
  elif a <= -180:
    a += 360
  return a

def compare_hit_anchor(query, hit):
  d = Data()
  # N anchor twist difference
  d.ntwist_diff = abs(norm_angle(query.Ntwist - hit.Ntwist))
  # C anchor twist difference
  d.ctwist_diff = abs(norm_angle(query.Ctwist - hit.Ctwist))
  return d

def find_contacts(gh, residue, maxdist=6.0):
  atoms = gh.atoms

  # gh is a GeometricHash object
  # atoms is a list of Atom objects that corresponds to the co-ordinates used to initialise gh
  # residue is a Residue object to find contacts for
  if maxdist <= 0:
    return []

  a = residue.CB
  if a is None:
    a = residue.CA
  if a is None:
    return []

  results = []
  nbrs = gh.get_neighbours(a.xyz, 999, max_dist=maxdist)
  for n in nbrs:
    ix, dst = n
    a = atoms[ix]

    element = a.atom[0:1]
    if element.isdigit():
      element = a.atom[1:2]
    if element == "H" or element.isdigit() or not element:
      continue

    if a.chain != residue.chain or abs(a.ires - residue.ires) > 10:
      results.append(a)

  return results

def find_contacts_simple(residues, i, maxdist=6.0):
  maxsqrdist = maxdist ** 2
  a1 = residues[i].CB
  if a1 is None:
    a1 = residues[i].CA
  if a1 is None:
    return []

  results = []
  #for residue in residues[:i-10] + residues[i+11:]:
  for residue in residues:
    a2 = residue.CB
    if a2 is None:
      a2 = residue.CA
    if a2 is None:
      continue

    if a1.chain != a2.chain or abs(a1.ires - a2.ires) > 10:
      sqd = sqrdist(a1, a2)
      if sqd < maxsqrdist:
        results.append(a2)

  return results

def get_contact_class(thischain, contact_atoms):
  x = 0 # no contacts
  for a in contact_atoms:
    if a.chain != thischain:
      x |= 1 # inter-chain contact
    else:
      x |= 2 # intra-chain contact

  # x = 0 ... no contacts
  # x = 1 ... only inter-chain contacts
  # x = 2 ... only intra-chain contacts
  # x = 3 ... both intra and inter-chain contacts
  return str(x)

def make_contact_gh(protein):
  coords = []
  gh_atoms = []
  for chain in protein:
    for r in chain:
      a = r.CB
      if a is None:
        a = r.CA
      if a is None:
        continue
      coords.append(a.xyz)
      gh_atoms.append(a)

  for chain in protein.ligands:
    for r in chain:
      for a in r:
        element = a.atom[0:1]
        if element.isdigit():
          element = a.atom[1:2]
        if element == "H" or element.isdigit() or not element:
          continue
        coords.append(a.xyz)
        gh_atoms.append(a)

  gh = GeometricHash(numpy.array(coords))
  gh.atoms = gh_atoms

  return gh


"""
Experimental code related to cifar
"""
import numpy as np
import ubelt as ub
import torch
# import torchvision
import pandas as pd
from torchvision.datasets import cifar
from netharn import xpu_device
from netharn import monitor
from netharn import initializers
from netharn import hyperparams
from netharn import fit_harness
from netharn.transforms import (ImageCenterScale,)
# from netharn.transforms import (RandomWarpAffine, RandomGamma, RandomBlur,)
import imgaug as ia
import imgaug.augmenters as iaa
from netharn import util
import netharn as nh

class CropTo(iaa.Augmenter):
    def __init__(self, shape,  name=None, deterministic=False, random_state=None):
        super(CropTo, self).__init__(name=name, deterministic=deterministic, random_state=random_state)
        self.shape = shape

    def _augment_images(self, images, random_state, parents, hooks):
        result = []
        nb_images = len(images)
        seeds = random_state.randint(0, 10**6, (nb_images,))
        for i in range(nb_images):
            seed = seeds[i]
            height, width = images[i].shape[0:2]
            top, bot, left, right = self._draw_samples_image(seed, height, width)

            image_cr = images[i][top:bot, left:right]
            image_cr = np.pad(image_cr, ((1, 1), (1, 1)), mode='constant')

            result.append(image_cr)
        return result

    def _augment_keypoints(self, keypoints_on_images, random_state, parents, hooks):
        result = []
        nb_images = len(keypoints_on_images)
        seeds = random_state.randint(0, 10**6, (nb_images,))
        for i, keypoints_on_image in enumerate(keypoints_on_images):
            seed = seeds[i]
            height, width = keypoints_on_image.shape[0:2]
            top, bot, left, right = self._draw_samples_image(seed, height, width)
            shifted = keypoints_on_image.shift(x=-left, y=-top)
            shifted.shape = (
                height - top - bot,
                width - left - right
            ) + shifted.shape[2:]
            result.append(shifted)
        return result

    def _draw_samples_image(self, seed, height, width):
        """
        height = 32
        width = 32
        h, w = shape = (30, 30)
        random_state = np.random
        """
        random_state = ia.new_random_state(seed)
        h, w = self.shape

        assert w <= width, '{} {}'.format(w, width)
        assert h <= height, '{} {}'.format(h, height)
        space_h = height - h
        space_w = width - w

        top = random_state.randint(0, space_h + 1)
        bot = height - (space_h - top)

        left = random_state.randint(0, space_w + 1)
        right = width - (space_w - left)

        sub = [top, bot, left, right]
        return sub

    def get_parameters(self):
        return [self.shape]

class Task(object):
    def __init__(task, labelnames=None, ignore_labelnames=[], alias={}):
        if labelnames is not None:
            task.set_labelnames(labelnames, ignore_labelnames, alias)

    def set_labelnames(task, labelnames, ignore_labelnames=[], alias={}):
        task.labelnames = list(labelnames)
        task.labelname_alias = alias
        task.ignore_labelnames = ignore_labelnames

        # Remove aliased classes
        for k in alias.keys():
            if k in task.labelnames:
                task.labelnames.remove(k)

        # Assign an integer label to each labelname
        task.labelname_to_id = ub.invert_dict(dict(enumerate(task.labelnames)))

        # Map aliased classes to a different label
        for k, v in alias.items():
            task.labelname_to_id[k] = task.labelname_to_id[v]

        task.ignore_labelnames = ignore_labelnames
        task.ignore_labels = np.array(
            list(ub.take(task.labelname_to_id, task.ignore_labelnames)))

        task.labels = np.arange(len(task.labelnames))
        task.relevant_labels = np.setdiff1d(task.labels, task.ignore_labels)

def radial_fourier_mask(img_chw, radius=11, axis=None, clip=None):
    """
    In [1] they use a radius of 11.0 on CIFAR-10.

    Args:
        img_chw (ndarray): assumed to be float 01

    References:
        [1] Jo and Bengio "Measuring the tendency of CNNs to Learn Surface Statistical Regularities" 2017.
        https://docs.opencv.org/3.0-beta/doc/py_tutorials/py_imgproc/py_transforms/py_fourier_transform/py_fourier_transform.html

    CommandLine:
        python examples/cifar.py radial_fourier_mask --show

    Example:
        >>> dset = cifar_training_datasets()['test']
        >>> dset.center_inputs = None
        >>> img_tensor, label = dset[7]
        >>> img_chw = img_tensor.numpy()
        >>> out = radial_fourier_mask(img_chw, radius=11)
        >>> # xdoc: REQUIRES(--show)
        >>> nh.util.qtensure()
        >>> def keepdim(func):
        >>>     def _wrap(im):
        >>>         needs_transpose = (im.shape[0] == 3)
        >>>         if needs_transpose:
        >>>             im = im.transpose(1, 2, 0)
        >>>         out = func(im)
        >>>         if needs_transpose:
        >>>             out = out.transpose(2, 0, 1)
        >>>         return out
        >>>     return _wrap
        >>> @keepdim
        >>> def rgb_to_bgr(im):
        >>>     return util.convert_colorspace(im, src_space='rgb', dst_space='bgr')
        >>> @keepdim
        >>> def bgr_to_lab(im):
        >>>     return util.convert_colorspace(im, src_space='bgr', dst_space='lab')
        >>> @keepdim
        >>> def lab_to_bgr(im):
        >>>     return util.convert_colorspace(im, src_space='lab', dst_space='bgr')
        >>> @keepdim
        >>> def bgr_to_yuv(im):
        >>>     return util.convert_colorspace(im, src_space='bgr', dst_space='yuv')
        >>> @keepdim
        >>> def yuv_to_bgr(im):
        >>>     return util.convert_colorspace(im, src_space='yuv', dst_space='bgr')
        >>> dpath = ub.ensuredir('./fouriertest')
        >>> from matplotlib import pyplot as plt
        >>> for x in ub.ProgIter(range(100)):
        >>>     img_tensor, label = dset[x]
        >>>     img_chw = img_tensor.numpy()
        >>>     bgr_img = rgb_to_bgr(img_chw)
        >>>     nh.util.imshow(bgr_img.transpose(1, 2, 0), fnum=1)
        >>>     pnum_ = nh.util.PlotNums(nRows=4, nCols=5)
        >>>     for r in range(0, 17):
        >>>         imgt = radial_fourier_mask(bgr_img, r, clip=(0, 1))
        >>>         nh.util.imshow(imgt.transpose(1, 2, 0), pnum=pnum_(), fnum=2)
        >>>         plt.gca().set_title('r = {}'.format(r))
        >>>     nh.util.set_figtitle('BGR')
        >>>     plt.gcf().savefig(join(dpath, '{}_{:08d}.png'.format('bgr', x)))
        >>>     pnum_ = nh.util.PlotNums(nRows=4, nCols=5)
        >>>     for r in range(0, 17):
        >>>         imgt = lab_to_bgr(radial_fourier_mask(bgr_to_lab(bgr_img), r)).transpose(1, 2, 0)
        >>>         nh.util.imshow(imgt, pnum=pnum_(), fnum=3)
        >>>         plt.gca().set_title('r = {}'.format(r))
        >>>         #imgt = lab_to_bgr(to_lab(bgr_img)).transpose(1, 2, 0)
        >>>         #nh.util.imshow(lab_to_bgr(to_lab(bgr_img)).transpose(1, 2, 0), pnum=pnum_(), fnum=2)
        >>>     nh.util.set_figtitle('LAB')
        >>>     plt.gcf().savefig(join(dpath, '{}_{:08d}.png'.format('lab', x)))
        >>>     pnum_ = nh.util.PlotNums(nRows=4, nCols=5)
        >>>     for r in range(0, 17):
        >>>         imgt = yuv_to_bgr(radial_fourier_mask(bgr_to_yuv(bgr_img), r, clip=(0., 1.))).transpose(1, 2, 0)
        >>>         nh.util.imshow(imgt, pnum=pnum_(), fnum=4)
        >>>         plt.gca().set_title('r = {}'.format(r))
        >>>     nh.util.set_figtitle('YUV')
        >>>     plt.gcf().savefig(join(dpath, '{}_{:08d}.png'.format('yuv', x)))
        >>> nh.util.show_if_requested()

    Ignore:
        im_chw = bgr_to_lab(bgr_img)
    """
    import cv2
    rows, cols = img_chw.shape[1:3]

    def fourier(s):
        # note: cv2 functions would probably be faster here
        return np.fft.fftshift(np.fft.fft2(s))

    def inv_fourier(f):
        # use real because LAB has negative components
        return np.real(np.fft.ifft2(np.fft.ifftshift(f)))

    diam = radius * 2
    left = int(np.floor((cols - diam) / 2))
    right = int(np.ceil((cols - diam) / 2))
    top = int(np.floor((rows - diam) / 2))
    bot = int(np.ceil((rows - diam) / 2))

    # element = skimage.morphology.disk(radius)
    # mask = np.pad(element, ((top, bot), (left, right)), 'constant')
    if diam > 0:
        element = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (diam, diam))
        mask = cv2.copyMakeBorder(element, top, bot, left, right, cv2.BORDER_CONSTANT, value=0)
    else:
        mask = 0

    out = np.empty_like(img_chw)
    if axis is None:
        for i, s in enumerate(img_chw):
            # hadamard product (aka simple element-wise multiplication)
            out[i] = inv_fourier(fourier(s) * mask)
    else:
        for i, s in enumerate(img_chw):
            if i in axis:
                # hadamard product (aka simple element-wise multiplication)
                out[i] = inv_fourier(fourier(s) * mask)
            else:
                out[i] = s
    if clip:
        out = np.clip(out, *clip)
    return out

    # nrows = cv2.getOptimalDFTSize(rows)
    # ncols = cv2.getOptimalDFTSize(cols)
    # right = ncols - cols
    # bottom = nrows - rows
    # if right or bottom:
    #     bordertype = cv2.BORDER_CONSTANT  # just to avoid line breakup in PDF file
    #     nimg = cv2.copyMakeBorder(img, 0, bottom, 0, right, bordertype, value=0)
    # dft_chans = [cv2.dft(chan, flags=cv2.DFT_COMPLEX_OUTPUT) for chan in img_chw]
    # dft = np.dstack(dft_chans).transpose(2, 0, 1)
    # dft_mag = np.dstack([(c ** 2).sum(axis=-1) for c in dft_chans]).transpose(2, 0, 1)
    # dft_shift = [np.fft.fftshift(c) for c in dft_chans]
    # dft_mag = np.dstack([(c ** 2).sum(axis=-1) for c in dft_shift]).transpose(2, 0, 1)
    # dft_filt_shift = [c * mask[:, :, None] for c in dft_shift]
    # dft_filt = [np.fft.ifftshift(c) for c in dft_filt_shift]
    # idft_filt = [cv2.idft(c) for c in dft_filt]
    # img_filt = np.dstack([np.linalg.norm(c, axis=-1) for c in idft_filt])
    # nh.util.imshow(dft_mag.transpose(1, 2, 0), norm=True)
    # if False:
    #     nh.util.imshow(np.log(dft_mag[0]), norm=True, pnum=(1, 3, 1))
    #     nh.util.imshow(np.log(dft_mag[1]), norm=True, pnum=(1, 3, 2))
    #     nh.util.imshow(np.log(dft_mag[2]), norm=True, pnum=(1, 3, 3))

def zca_whitening_matrix(X):
    """
    Function to compute ZCA whitening matrix (aka Mahalanobis whitening).
    Args:
        X (ndarray): [M x N] matrix, Rows: Variables, Columns: Observations

    Returns:
        ZCAMatrix: [M x M] matrix

    References:
        https://stackoverflow.com/a/38590790/887074

    Example:
        >>> rng = np.random.RandomState(0)
        >>> # Construct a matrix of observations from grayscale 8x8 images
        >>> gray_images = [rng.rand(8, 8) for _ in range(1000)]
        >>> X = np.array([img.ravel() for img in gray_images]).T
        >>> M = zca_whitening_matrix(X)
        >>> img = gray_images[0]
        >>> norm = M.dot(img.ravel()).reshape(8, 8)
        >>> # ... for the RGB channels of color images
        >>> rgb_images = [rng.rand(3, 8, 8) for _ in range(1000)]
        >>> #X = np.array([img.mean(axis=(1, 2)) for img in rgb_images]).T
        >>> X = np.hstack([img.reshape(3, -1) for img in rgb_images])
        >>> M = zca_whitening_matrix(X)
        >>> img = rgb_images[0]
        >>> norm = M.dot(img.reshape(3, 64)).reshape(3, 8, 8)
    """
    # Covariance matrix [column-wise variables]: Sigma = (X-mu)' * (X-mu) / N
    sigma = np.cov(X, rowvar=True)  # [M x M]
    # Singular Value Decomposition. X = U * np.diag(S) * V
    U, S, V = np.linalg.svd(sigma)
    # U: [M x M] eigenvectors of sigma.
    # S: [M x 1] eigenvalues of sigma.
    # V: [M x M] transpose of U
    # Whitening constant: prevents division by zero
    epsilon = 1e-5
    L = np.diag(1.0 / np.sqrt(S + epsilon))
    # ZCA Whitening matrix: U * Lambda * U'
    ZCAMatrix = np.dot(U, np.dot(L, U.T))  # [M x M]
    return ZCAMatrix

class CIFAR10_Task(Task):
    """
    task = CIFAR10_Task()
    task._initialize()
    ignore_labelnames = []
    alias = {}
    """
    def __init__(task, root=None):
        if root is None:
            root = ub.ensure_app_cache_dir('netharn')
        task.root = root
        task._initialize()

    def _initialize(task):
        from os.path import join
        import pickle
        train_dset = cifar.CIFAR10(root=task.root, download=False, train=True)

        fpath = join(train_dset.root,
                     cifar.CIFAR10.base_folder, 'batches.meta')
        with open(fpath, 'rb') as fo:
            entry = pickle.load(fo, encoding='latin1')
            labelnames = entry['label_names']
        task.set_labelnames(labelnames)

class CIFAR100_Task(Task):
    """
    task = CIFAR100_Task()
    task._initialize()
    ignore_labelnames = []
    alias = {}
    """
    def __init__(task, root=None):
        if root is None:
            root = ub.ensure_app_cache_dir('netharn')
        task.root = root
        task._initialize()

    def _initialize(task):
        from os.path import join
        import pickle
        train_dset = cifar.CIFAR100(root=task.root, download=False, train=True)

        fpath = join(train_dset.root, cifar.CIFAR100.base_folder, 'meta')
        with open(fpath, 'rb') as fo:
            entry = pickle.load(fo, encoding='latin1')
            labelnames = entry['fine_label_names']
        task.set_labelnames(labelnames)

def mutex_clf_gt_info(gt_labels, task):
    """
    gt_labels = train_dset.train_labels
    """
    index = pd.Index(task.labels, name='label')
    gtstats = pd.DataFrame(0, index=index, columns=['freq'], dtype=np.int)

    label_freq = pd.value_counts(gt_labels)
    gtstats.freq = pd.to_numeric(label_freq)

    gtstats['classname'] = list(ub.take(task.labelnames, gtstats.index))
    gtstats['mf_weight'] = gtstats.freq.median() / gtstats.freq
    gtstats.loc[~np.isfinite(gtstats.mf_weight), 'mf_weight'] = 1

    # Clip weights, so nothing gets crazy high weights, low weights are ok
    gtstats = gtstats.sort_index()
    gtstats.index.name = 'label'
    gtstats = gtstats.reset_index().set_index('classname', drop=False)
    return gtstats

class InMemoryInputs(ub.NiceRepr):
    """
    Change inputs.Inputs to OnDiskInputs
    """
    def __init__(inputs, tag=''):
        inputs.tag = tag
        inputs.im = None
        inputs.gt = None
        inputs.colorspace = None
        inputs.input_id = None

    def __nice__(inputs):
        n = len(inputs)
        return '{} {}'.format(inputs.tag, n)

    def __len__(inputs):
        if inputs.im is not None:
            n = len(inputs.im)
        elif inputs.gt is not None:
            n = len(inputs.gt)
        else:
            n = 0
        return n

    @classmethod
    def from_bhwc_rgb(cls, bhwc, labels=None, **kw):
        # convert to bhwc
        inputs = cls(**kw)
        inputs.im = bhwc
        inputs.gt = labels
        inputs.colorspace = 'rgb'
        return inputs

    def convert_colorspace(inputs, colorspace, inplace=False):
        if colorspace.lower() == inputs.colorspace.lower():
            if not inplace:
                return inputs.im
            return
        im_out = np.empty_like(inputs.im)
        dst = np.ascontiguousarray(np.empty_like(inputs.im[0]))
        for ix, im in enumerate(inputs.im):
            util.convert_colorspace(im, src_space=inputs.colorspace,
                                    dst_space=colorspace, dst=dst)
            im_out[ix] = dst
        if inplace:
            inputs.im = im_out
            inputs.colorspace = colorspace
        else:
            return im_out

    def take(inputs, idxs, **kw):
        new_inputs = inputs.__class__(**kw)
        new_inputs.im = inputs.im.take(idxs, axis=0)
        new_inputs.gt = inputs.gt.take(idxs, axis=0)
        new_inputs.colorspace = inputs.colorspace
        return new_inputs

    def prepare_id(self, force=False):
        if self.input_id is not None and not force:
            return

        depends = []
        depends.append(self.im)
        depends.append(self.gt)

    def _set_id_from_dependency(self, depends):
        """
        Allow for arbitrary representation of dependencies
        (user must ensure that it is consistent)
        """
        print('Preparing id for {} images'.format(self.tag))
        abbrev = 8
        hashid = util.hash_data(depends)[:abbrev]
        n_input = len(self)
        self.input_id = '{}-{}'.format(n_input, hashid)
        print(' * n_input = {}'.format(n_input))
        print(' * input_id = {}'.format(self.input_id))

class CIFAR_Wrapper(torch.utils.data.Dataset):  # cifar.CIFAR10):
    def __init__(dset, inputs, task, workdir, output_colorspace='RGB'):
        dset.inputs = inputs
        dset.task = task

        dset.output_colorspace = output_colorspace

        dset.rng = np.random.RandomState(432432)

        inputs_base = ub.ensuredir((workdir, 'inputs'))
        inputs.base_dpath = inputs_base
        if len(inputs):
            inputs.prepare_id()
            dset.input_id = inputs.input_id
            dset.with_gt = dset.inputs.gt is not None
        else:
            dset.input_id = ''

        # TODO: only use horizontal flipping and translation by 4 pixels to
        # match results from other papers
        # https://arxiv.org/pdf/1603.09382.pdf page 8

        dset.augment = None
        # dset.im_augment = torchvision.transforms.Compose([
        #     RandomGamma(rng=dset.rng),
        #     RandomBlur(rng=dset.rng),
        # ])
        # dset.rand_aff = RandomWarpAffine(dset.rng)

        augmentors = [
            # iaa.Sometimes(.8, iaa.ContrastNormalization((0.2, 1.8))),
            iaa.Fliplr(p=.5),
            iaa.Affine(translate_px={'x': (-1, 1), 'y': (-1, 1)}),

            # CropTo((30, 30)),
            # iaa.Crop(px=(1, 1, 1, 1)),
            # imgaug.Brightness(63),
            # imgaug.RandomCrop((30, 30)),
            # imgaug.MeanVarianceNormalize(all_channel=True)
        ]
        dset.augmenter = iaa.Sequential(augmentors)
        # iaa.Sequential([
        #     iaa.Affine(translate_px={"x":-40}),
        #     iaa.AdditiveGaussianNoise(scale=0.1*255)
        # ])

        # dset.rand_aff = RandomWarpAffine(
        #     dset.rng, tx_pdf=(-2, 2), ty_pdf=(-2, 2), flip_lr_prob=.5,
        #     zoom_pdf=None, shear_pdf=None, flip_ud_prob=None,
        #     enable_stretch=None, default_distribution='uniform')

        dset.center_inputs = None

    def _make_normalizer(dset, mode='independent'):
        """
        Example:
            >>> inputs, task = cifar_inputs(train=True)
            >>> workdir = ub.ensuredir(ub.truepath('~/data/work/cifar'))
            >>> dset = CIFAR_Wrapper(inputs, task, workdir, 'RGB')
            >>> center_inputs = dset._make_normalizer('independent')
        """
        if len(dset.inputs):
            # compute normalizers in the output colorspace
            out_im = dset.inputs.convert_colorspace(dset.output_colorspace,
                                                    inplace=False)
            if mode == 'dependant':
                # dependent centering per channel (for RGB)
                im_mean = out_im.mean()
                im_scale = out_im.std()
            elif mode == 'independent':
                # Independent centering per channel (for LAB)
                im_mean = out_im.mean(axis=(0, 1, 2))
                im_scale = out_im.std(axis=(0, 1, 2))

            center_inputs = ImageCenterScale(im_mean, im_scale)

        dset.center_inputs = center_inputs
        return center_inputs

    def __len__(dset):
        return len(dset.inputs)

    def load_inputs(dset, index):
        """
        Ignore:
            >>> inputs, task = cifar_inputs(train=False)
            >>> workdir = ub.ensuredir(ub.truepath('~/data/work/cifar'))
            >>> dset = CIFAR_Wrapper(inputs, task, workdir, 'LAB')
            >>> dset._make_normalizer('independent')
            >>> index = 0
            >>> im, gt = dset.load_inputs(index)

        Example:
            >>> inputs, task = cifar_inputs(train=False)
            >>> workdir = ub.ensuredir(ub.truepath('~/data/work/cifar'))
            >>> dset = CIFAR_Wrapper(inputs, task, workdir, 'RGB')
            >>> index = 0
            >>> im, gt = dset.load_inputs(index)
            >>> from netharn.util import mplutil
            >>> mplutil.qtensure()
            >>> dset = CIFAR_Wrapper(inputs, task, workdir, 'RGB')
            >>> dset.augment = True
            >>> im, gt = dset.load_inputs(index)
            >>> mplutil.imshow(im, colorspace='rgb')

            >>> dset = CIFAR_Wrapper(inputs, task, workdir, 'LAB')
            >>> dset.augment = True
            >>> im, gt = dset.load_inputs(index)
            >>> mplutil.imshow(im, colorspace='LAB')
        """
        assert dset.inputs.colorspace.lower() == 'rgb', (
            'we must be in rgb for augmentation')
        im = dset.inputs.im[index]

        if dset.inputs.gt is not None:
            gt = dset.inputs.gt[index]
        else:
            gt = None

        if dset.augment:
            # Image augmentation must be done in RGB
            # Augment intensity independently
            # im = dset.im_augment(im)
            # Augment geometry consistently

            # params = dset.rand_aff.random_params()
            # im = dset.rand_aff.warp(im, params, interp='cubic', backend='cv2')

            im = util.convert_colorspace(im, src_space=dset.inputs.colorspace,
                                         dst_space='rgb')
            # Do augmentation in uint8 RGB
            im = (im * 255).astype(np.uint8)
            im = dset.augmenter.augment_image(im)
            im = (im / 255).astype(np.float32)
            im = util.convert_colorspace(im, src_space='rgb',
                                         dst_space=dset.output_colorspace)
        else:
            im = util.convert_colorspace(im, src_space=dset.inputs.colorspace,
                                         dst_space=dset.output_colorspace)
        # Do centering of inputs
        if dset.center_inputs:
            im = dset.center_inputs(im)
        return im, gt

    def __getitem__(dset, index):
        from netharn import im_loaders
        im, gt = dset.load_inputs(index)
        input_tensor = im_loaders.numpy_image_to_float_tensor(im)

        if dset.with_gt:
            # print('gotitem: ' + str(data_tensor.shape))
            # print('gt_tensor: ' + str(gt_tensor.shape))
            return input_tensor, gt
        else:
            return input_tensor

    @property
    def n_channels(dset):
        return 3

    @property
    def n_classes(dset):
        return int(dset.task.labels.max() + 1)

    @property
    def ignore_labels(dset):
        return dset.task.ignore_labels

    def class_weights(dset):
        """
            >>> from netharn.live.sseg_train import *
            >>> dset = load_task_dataset('urban_mapper_3d')['train']
            >>> dset.class_weights()
        """
        # # Handle class weights
        # print('prep class weights')
        # gtstats = dset.inputs.prepare_gtstats(dset.task)
        # gtstats = dset.inputs.gtstats
        # # Take class weights (ensure they are in the same order as labels)
        # mfweight_dict = gtstats['mf_weight'].to_dict()
        # class_weights = np.array(list(ub.take(mfweight_dict, dset.task.classnames)))
        # class_weights[dset.task.ignore_labels] = 0
        # # HACK
        # # class_weights[0] = 1.0
        # # class_weights[1] = 0.7
        # print('class_weights = {!r}'.format(class_weights))
        # print('class_names   = {!r}'.format(dset.task.classnames))
        class_weights = np.ones(dset.n_classes)
        return class_weights

def cifar_inputs(train=False, cifar_num=10):
    root = ub.ensure_app_cache_dir('netharn')

    if cifar_num == 10:
        train_dset = cifar.CIFAR10(root=root, download=True, train=train)
        task = CIFAR10_Task()
    else:
        train_dset = cifar.CIFAR100(root=root, download=True, train=train)
        task = CIFAR100_Task()
    if train:
        bchw = (train_dset.train_data).astype(np.float32) / 255.0
        labels = np.array(train_dset.train_labels)
    else:
        bchw = (train_dset.test_data).astype(np.float32) / 255.0
        labels = np.array(train_dset.test_labels)
    inputs = InMemoryInputs.from_bhwc_rgb(bchw, labels=labels)
    if train:
        inputs.tag = 'learn'
    else:
        inputs.tag = 'test'
    return inputs, task

def cifar_training_datasets(output_colorspace='RGB', norm_mode='independent',
                            cifar_num=10):
    """
    Example:
        >>> datasets = cifar_training_datasets()
    """
    inputs, task = cifar_inputs(train=True, cifar_num=cifar_num)

    # split training into train / validation
    # 45K / 5K validation split was used in densenet and resnet papers.
    # https://arxiv.org/pdf/1512.03385.pdf page 7
    # https://arxiv.org/pdf/1608.06993.pdf page 5

    vali_frac = .1  # 10%  is 5K images
    n_vali = int(len(inputs) * vali_frac)
    # n_vali = 10000  # 10K validation as in http://torch.ch/blog/2015/07/30/cifar.html

    # the gt indexes seem to already be scrambled, I think other papers sample
    # validation from the end, so lets do that
    # The NIN paper https://arxiv.org/pdf/1312.4400.pdf in section 4 mentions
    # that it uses the last 10K images for validation
    input_idxs = np.arange(len(inputs))
    # or just uncomment this line for reproducable random sampling
    # input_idxs = util.random_indices(len(inputs), seed=1184576173)

    train_idxs = sorted(input_idxs[:-n_vali])
    vali_idxs = sorted(input_idxs[-n_vali:])

    train_inputs = inputs.take(train_idxs, tag='train')
    vali_inputs = inputs.take(vali_idxs, tag='vali')
    test_inputs, _ = cifar_inputs(train=False, cifar_num=cifar_num)
    # The dataset name and indices should fully specifiy dependencies
    train_inputs._set_id_from_dependency(
        ['cifar{}-train'.format(cifar_num), train_idxs])
    vali_inputs._set_id_from_dependency(
        ['cifar{}-train'.format(cifar_num), vali_idxs])
    test_inputs._set_id_from_dependency(['cifar{}-test'.format(cifar_num)])

    workdir = ub.ensuredir(ub.truepath('~/data/work/cifar'))

    train_dset = CIFAR_Wrapper(
        train_inputs, task, workdir, output_colorspace=output_colorspace)
    vali_dset = CIFAR_Wrapper(
        vali_inputs, task, workdir, output_colorspace=output_colorspace)
    test_dset = CIFAR_Wrapper(test_inputs, task, workdir,
                              output_colorspace=output_colorspace)
    print('built datasets')

    datasets = {
        'train': train_dset,
        'vali': vali_dset,
        'test': test_dset,
    }

    print('computing normalizers')
    datasets['train'].center_inputs = datasets['train']._make_normalizer(
        norm_mode)
    for key in datasets.keys():
        datasets[key].center_inputs = datasets['train'].center_inputs
    print('computed normalizers')

    datasets['train'].augment = True
    return datasets

def train():
    """
    Example:
        >>> train()
    """
    import random
    np.random.seed(1031726816 % 4294967295)
    torch.manual_seed(137852547 % 4294967295)
    random.seed(2497950049 % 4294967295)

    xpu = xpu_device.XPU.from_argv()
    print('Chosen xpu = {!r}'.format(xpu))

    cifar_num = 10

    if ub.argflag('--lab'):
        datasets = cifar_training_datasets(
            output_colorspace='LAB', norm_mode='independent', cifar_num=cifar_num)
    elif ub.argflag('--rgb'):
        datasets = cifar_training_datasets(
            output_colorspace='RGB', norm_mode='independent', cifar_num=cifar_num)
    elif ub.argflag('--rgb-dep'):
        datasets = cifar_training_datasets(
            output_colorspace='RGB', norm_mode='dependant', cifar_num=cifar_num)
    else:
        raise AssertionError('specify --rgb / --lab')

    import netharn.models.densenet

    # batch_size = (128 // 3) * 3
    batch_size = 64

    # initializer_ = (initializers.KaimingNormal, {
    #     'nonlinearity': 'relu',
    # })

    lr = 0.1
    initializer_ = (initializers.LSUV, {})

    hyper = hyperparams.HyperParams(
        model=(netharn.models.densenet.DenseNet, {
            'cifar': True,
            'block_config': (32, 32, 32),  # 100 layer depth
            'num_classes': datasets['train'].n_classes,
            'drop_rate': float(ub.argval('--drop_rate', default=.2)),
            'groups': 1,
        }),
        optimizer=(torch.optim.SGD, {
            # 'weight_decay': .0005,
            'weight_decay': float(ub.argval('--weight_decay', default=.0005)),
            'momentum': 0.9,
            'nesterov': True,
            'lr': 0.1,
        }),
        scheduler=(nh.schedulers.ListedLR, {
            'points': {
                0: lr,
                150: lr * 0.1,
                250: lr * 0.01,
            },
            'interpolate': False
        }),
        monitor=(nh.Monitor, {
            'minimize': ['loss'],
            'maximize': ['mAP'],
            'patience': 314,
            'max_epoch': 314,
        }),
        initializer=initializer_,
        criterion=(torch.nn.CrossEntropyLoss, {
        }),
        # Specify anything else that is special about your hyperparams here
        # Especially if you make a custom_batch_runner
        augment=str(datasets['train'].augmenter),
        other=ub.dict_union({
            # TODO: type of augmentation as a parameter dependency
            # 'augmenter': str(datasets['train'].augmenter),
            # 'augment': datasets['train'].augment,
            'batch_size': batch_size,
            'colorspace': datasets['train'].output_colorspace,
            'n_classes': datasets['train'].n_classes,
            # 'center_inputs': datasets['train'].center_inputs,
        }, datasets['train'].center_inputs.__dict__),
    )
    # if ub.argflag('--rgb-indie'):
    #     hyper.other['norm'] = 'dependant'
    hyper.input_ids['train'] = datasets['train'].input_id

    xpu = xpu_device.XPU.cast('auto')
    print('xpu = {}'.format(xpu))

    data_kw = {'batch_size': batch_size}
    if xpu.is_gpu():
        data_kw.update({'num_workers': 8, 'pin_memory': True})

    tags = ['train', 'vali', 'test']

    loaders = ub.odict()
    for tag in tags:
        dset = datasets[tag]
        shuffle = tag == 'train'
        data_kw_ = data_kw.copy()
        if tag != 'train':
            data_kw_['batch_size'] = max(batch_size // 4, 1)
        loader = torch.utils.data.DataLoader(dset, shuffle=shuffle, **data_kw_)
        loaders[tag] = loader

    harn = fit_harness.FitHarness(
        hyper=hyper, datasets=datasets, xpu=xpu,
        loaders=loaders,
    )
    # harn.monitor = early_stop.EarlyStop(patience=40)
    harn.monitor = monitor.Monitor(min_keys=['loss'],
                                   max_keys=['global_acc', 'class_acc'],
                                   patience=40)

    # ignore_label = datasets['train'].ignore_label
    # from netharn import metrics

    workdir = ub.ensuredir('train_cifar_work')
    harn.setup_dpath(workdir)

    harn.run()

 """
        python examples/cifar.py train --lab
        python examples/cifar.py train --rgb-indie
 """


from . import gwem_resampling_utils as sampler_functions
import pandas as pd
import numpy as np
import bilby
import argparse
import os

def main():

    parser = argparse.ArgumentParser(
        description="Inference on binary source parameters with kilonova ejecta posterior and GW source posterior given."
    )
    parser.add_argument("--outdir", metavar="PATH", type=str, required=True)
    parser.add_argument(
            "--GWsamples",
            metavar="PATH",
            type=str,
            required=True,
            help="If no posterior files are available, use gwsamples_creation.py to generate dummy GWsamples."
    )
    parser.add_argument(
            "--EMsamples",
            metavar="PATH",
            type=str,
            required=True,
            help="posterior samples file from a previous Bayesian inference run on EM signals (e.g. Kilonova inference or Kilonova+GRB inference.")
    parser.add_argument(
            "--EOSpath",
            metavar="PATH",
            type=str,
            required=True,
            help="Path of EOS folder, e.g. 15nsat_cse_uniform_R14 (located: https://zenodo.org/record/6106130#.YoysIHVBwUG)"
    )
    parser.add_argument(
            "--Neos",
            metavar="Neos",
            type=int,
            required=True,
            help="Number of EOS files used for the inference."
    )
    parser.add_argument(
            "--nlive",
            metavar="nlive",
            type=int,
            required=False,
            default=1024
    )
    parser.add_argument(
        "--GWprior",
        metavar="PATH",
        type=str,
        required=True,
        help="Prior file used for the GW analysis",
    )
    parser.add_argument(
        "--EMprior",
        metavar="PATH",
        type=str,
        required=True,
        help="Prior file used for the EM eos analysis",
    )
    parser.add_argument(
        "--total-ejecta-mass",
        action="store_true",
        help="To run with total ejecta mass, if not activated, the two ejecta are consider seperately",
    )
    args = parser.parse_args()

    # read the GW samples
    GWsamples = pd.read_csv(args.GWsamples, header=0, delimiter=" ")
    # down sample
    weights = np.ones(len(GWsamples))
    weights /= np.sum(weights)
    GWsamples = GWsamples.sample(
        frac=30000 / len(GWsamples), weights=weights, random_state=42
    )

    # read the EM samples
    EMsamples = pd.read_csv(args.EMsamples, header=0, delimiter=" ")

    # read the prior files
    GWprior = bilby.gw.prior.PriorDict(args.GWprior)
    EMprior = bilby.gw.prior.PriorDict(args.EMprior)

    try:
        os.makedirs(args.outdir + "/pm/")
    except Exception:
        pass
    pymulti_kwargs = dict(
        outputfiles_basename=args.outdir + "/pm/",
        n_dims=5,
        n_live_points=args.nlive,
        verbose=True,
        resume=True,
        seed=42,
        importance_nested_sampling=False,
    )

    if args.total_ejecta_mass:
        solution = sampler_functions.TotalEjectaMassInference(
            GWsamples,
            EMsamples,
            GWprior,
            EMprior,
            args.Neos,
            args.EOSpath,
            **pymulti_kwargs
        )
    else:
        solution = sampler_functions.EjectaMassInference(
            GWsamples,
            EMsamples,
            GWprior,
            EMprior,
            args.Neos,
            args.EOSpath,
            **pymulti_kwargs
        )

    samples = solution.samples.T
    posterior_samples = dict()
    posterior_samples["chirp_mass"] = samples[0]
    posterior_samples["mass_ratio"] = samples[1]
    posterior_samples["EOS"] = samples[2]
    posterior_samples["alpha"] = samples[3]
    posterior_samples["zeta"] = samples[4]

    posterior_samples = pd.DataFrame.from_dict(posterior_samples)
    posterior_samples.to_csv(
        "{0}/posterior_samples.dat".format(args.outdir), sep=" ", index=False
    )

    sampler_functions.corner_plot(posterior_samples, solution, args.outdir)

if __name__ == "__main__":
    main()


# encoding=utf8
import abc
import copy
from typing import List, Any, Tuple

import numpy as np
from scipy.stats import norm

from smac.configspace import Configuration
from smac.configspace.util import convert_configurations_to_array
from smac.epm.base_epm import AbstractEPM
from smac.utils.logging import PickableLoggerAdapter

__author__ = "Aaron Klein, Marius Lindauer"
__copyright__ = "Copyright 2017, ML4AAD"
__license__ = "3-clause BSD"

class AbstractAcquisitionFunction(object, metaclass=abc.ABCMeta):
    """Abstract base class for acquisition function

    Attributes
    ----------
    model
    logger
    """

    def __init__(self, model: AbstractEPM):
        """Constructor

        Parameters
        ----------
        model : AbstractEPM
            Models the objective function.
        """
        self.model = model
        self._required_updates = ('model', )  # type: Tuple[str, ...]
        self.logger = PickableLoggerAdapter(self.__module__ + "." + self.__class__.__name__)

    def update(self, **kwargs: Any) -> None:
        """Update the acquisition function attributes required for calculation.

        This method will be called after fitting the model, but before maximizing the acquisition
        function. As an examples, EI uses it to update the current fmin.

        The default implementation only updates the attributes of the acqusition function which
        are already present.

        Parameters
        ----------
        kwargs
        """
        for key in self._required_updates:
            if key not in kwargs:
                raise ValueError(
                    'Acquisition function %s needs to be updated with key %s, but only got '
                    'keys %s.'
                    % (self.__class__.__name__, key, list(kwargs.keys()))
                )
        for key in kwargs:
            if key in self._required_updates:
                setattr(self, key, kwargs[key])

    def __call__(self, configurations: List[Configuration]) -> np.ndarray:
        """Computes the acquisition value for a given X

        Parameters
        ----------
        configurations : list
            The configurations where the acquisition function
            should be evaluated.

        Returns
        -------
        np.ndarray(N, 1)
            acquisition values for X
        """
        X = convert_configurations_to_array(configurations)
        if len(X.shape) == 1:
            X = X[np.newaxis, :]

        acq = self._compute(X)
        if np.any(np.isnan(acq)):
            idx = np.where(np.isnan(acq))[0]
            acq[idx, :] = -np.finfo(np.float).max
        return acq

    @abc.abstractmethod
    def _compute(self, X: np.ndarray) -> np.ndarray:
        """Computes the acquisition value for a given point X. This function has
        to be overwritten in a derived class.

        Parameters
        ----------
        X : np.ndarray
            The input points where the acquisition function
            should be evaluated. The dimensionality of X is (N, D), with N as
            the number of points to evaluate at and D is the number of
            dimensions of one X.

        Returns
        -------
        np.ndarray(N,1)
            Acquisition function values wrt X
        """
        raise NotImplementedError()

class IntegratedAcquisitionFunction(AbstractAcquisitionFunction):

    r"""Marginalize over Model hyperparameters to compute the integrated acquisition function.

    See "Practical Bayesian Optimization of Machine Learning Algorithms" by Jasper Snoek et al.
    (https://papers.nips.cc/paper/4522-practical-bayesian-optimization-of-machine-learning-algorithms.pdf)
    for further details.
    """

    def __init__(self, model: AbstractEPM, acquisition_function: AbstractAcquisitionFunction, **kwargs: Any):
        """Constructor

        Parameters
        ----------
        model : AbstractEPM
            The model needs to implement an additional attribute ``models`` which contains the different models to
            integrate over.
        kwargs
            Additional keyword arguments
        """

        super().__init__(model)
        self.long_name = 'Integrated Acquisition Function (%s)' % acquisition_function.__class__.__name__
        self.acq = acquisition_function
        self._functions = []  # type: List[AbstractAcquisitionFunction]
        self.eta = None

    def update(self, **kwargs: Any) -> None:
        """Update the acquisition functions values.

        This method will be called if the model is updated. E.g. entropy search uses it to update its approximation
        of P(x=x_min), EI uses it to update the current fmin.

        This implementation creates an acquisition function object for each model to integrate over and sets the
        respective attributes for each acquisition function object.

        Parameters
        ----------
        model : AbstractEPM
            The model needs to implement an additional attribute ``models`` which contains the different models to
            integrate over.
        kwargs
        """
        model = kwargs['model']
        del kwargs['model']
        if not hasattr(model, 'models') or len(model.models) == 0:
            raise ValueError('IntegratedAcquisitionFunction requires at least one model to integrate!')
        if len(self._functions) == 0 or len(self._functions) != len(model.models):
            self._functions = [copy.deepcopy(self.acq) for _ in model.models]
        for submodel, func in zip(model.models, self._functions):
            func.update(model=submodel, **kwargs)

    def _compute(self, X: np.ndarray) -> np.ndarray:
        """Computes the EI value and its derivatives.

        Parameters
        ----------
        X: np.ndarray(N, D), The input points where the acquisition function
            should be evaluated. The dimensionality of X is (N, D), with N as
            the number of points to evaluate at and D is the number of
            dimensions of one X.

        Returns
        -------
        np.ndarray(N,1)
            Expected Improvement of X
        """
        if self._functions is None:
            raise ValueError('Need to call update first!')
        return np.array([func._compute(X) for func in self._functions]).mean(axis=0)

class EI(AbstractAcquisitionFunction):

    r"""Computes for a given x the expected improvement as
    acquisition value.

    :math:`EI(X) := \mathbb{E}\left[ \max\{0, f(\mathbf{X^+}) - f_{t+1}(\mathbf{X}) - \xi \} \right]`,
    with :math:`f(X^+)` as the best location.
    """

    def __init__(self,
                 model: AbstractEPM,
                 par: float = 0.0):
        """Constructor

        Parameters
        ----------
        model : AbstractEPM
            A model that implements at least
                 - predict_marginalized_over_instances(X)
        par : float, default=0.0
            Controls the balance between exploration and exploitation of the
            acquisition function.
        """

        super(EI, self).__init__(model)
        self.long_name = 'Expected Improvement'
        self.par = par
        self.eta = None
        self._required_updates = ('model', 'eta')

    def _compute(self, X: np.ndarray) -> np.ndarray:
        """Computes the EI value and its derivatives.

        Parameters
        ----------
        X: np.ndarray(N, D), The input points where the acquisition function
            should be evaluated. The dimensionality of X is (N, D), with N as
            the number of points to evaluate at and D is the number of
            dimensions of one X.

        Returns
        -------
        np.ndarray(N,1)
            Expected Improvement of X
        """
        if len(X.shape) == 1:
            X = X[:, np.newaxis]

        m, v = self.model.predict_marginalized_over_instances(X)
        s = np.sqrt(v)

        if self.eta is None:
            raise ValueError('No current best specified. Call update('
                             'eta=<int>) to inform the acquisition function '
                             'about the current best value.')

        def calculate_f():
            z = (self.eta - m - self.par) / s
            return (self.eta - m - self.par) * norm.cdf(z) + s * norm.pdf(z)

        if np.any(s == 0.0):
            # if std is zero, we have observed x on all instances
            # using a RF, std should be never exactly 0.0
            # Avoid zero division by setting all zeros in s to one.
            # Consider the corresponding results in f to be zero.
            self.logger.warning("Predicted std is 0.0 for at least one sample.")
            s_copy = np.copy(s)
            s[s_copy == 0.0] = 1.0
            f = calculate_f()
            f[s_copy == 0.0] = 0.0
        else:
            f = calculate_f()
        if (f < 0).any():
            raise ValueError(
                "Expected Improvement is smaller than 0 for at least one "
                "sample.")

        return f

class EIPS(EI):
    def __init__(self,
                 model: AbstractEPM,
                 par: float = 0.0):
        r"""Computes for a given x the expected improvement as
        acquisition value.
        :math:`EI(X) := \frac{\mathbb{E}\left[\max\{0,f(\mathbf{X^+})-f_{t+1}(\mathbf{X})-\xi\right]\}]}{np.log(r(x))}`,
        with :math:`f(X^+)` as the best location and :math:`r(x)` as runtime.

        Parameters
        ----------
        model : AbstractEPM
            A model that implements at least
                 - predict_marginalized_over_instances(X) returning a tuples of
                   predicted cost and running time
        par : float, default=0.0
            Controls the balance between exploration and exploitation of the
            acquisition function.
        """
        super(EIPS, self).__init__(model, par=par)
        self.long_name = 'Expected Improvement per Second'

    def _compute(self, X: np.ndarray) -> np.ndarray:
        """Computes the EIPS value.

        Parameters
        ----------
        X: np.ndarray(N, D), The input point where the acquisition function
            should be evaluate. The dimensionality of X is (N, D), with N as
            the number of points to evaluate at and D is the number of
            dimensions of one X.

        Returns
        -------
        np.ndarray(N,1)
            Expected Improvement per Second of X
        """
        if len(X.shape) == 1:
            X = X[:, np.newaxis]

        m, v = self.model.predict_marginalized_over_instances(X)
        if m.shape[1] != 2:
            raise ValueError("m has wrong shape: %s != (-1, 2)" % str(m.shape))
        if v.shape[1] != 2:
            raise ValueError("v has wrong shape: %s != (-1, 2)" % str(v.shape))

        m_cost = m[:, 0]
        v_cost = v[:, 0]
        # The model already predicts log(runtime)
        m_runtime = m[:, 1]
        s = np.sqrt(v_cost)

        if self.eta is None:
            raise ValueError('No current best specified. Call update('
                             'eta=<int>) to inform the acquisition function '
                             'about the current best value.')

        def calculate_f():
            z = (self.eta - m_cost - self.par) / s
            f = (self.eta - m_cost - self.par) * norm.cdf(z) + s * norm.pdf(z)
            f = f / m_runtime
            return f

        if np.any(s == 0.0):
            # if std is zero, we have observed x on all instances
            # using a RF, std should be never exactly 0.0
            # Avoid zero division by setting all zeros in s to one.
            # Consider the corresponding results in f to be zero.
            self.logger.warning("Predicted std is 0.0 for at least one sample.")
            s_copy = np.copy(s)
            s[s_copy == 0.0] = 1.0
            f = calculate_f()
            f[s_copy == 0.0] = 0.0
        else:
            f = calculate_f()

        if (f < 0).any():
            raise ValueError(
                "Expected Improvement per Second is smaller than 0 "
                "for at least one sample.")

        return f.reshape((-1, 1))

class LogEI(AbstractAcquisitionFunction):

    def __init__(self,
                 model: AbstractEPM,
                 par: float = 0.0):
        r"""Computes for a given x the logarithm expected improvement as
        acquisition value.

        Parameters
        ----------
        model : AbstractEPM
            A model that implements at least
                 - predict_marginalized_over_instances(X)
        par : float, default=0.0
            Controls the balance between exploration and exploitation of the
            acquisition function.
        """
        super(LogEI, self).__init__(model)
        self.long_name = 'Expected Improvement'
        self.par = par
        self.eta = None
        self._required_updates = ('model', 'eta')

    def _compute(self, X: np.ndarray) -> np.ndarray:
        """Computes the EI value and its derivatives.

        Parameters
        ----------
        X: np.ndarray(N, D), The input points where the acquisition function
            should be evaluated. The dimensionality of X is (N, D), with N as
            the number of points to evaluate at and D is the number of
            dimensions of one X.

        Returns
        -------
        np.ndarray(N,1)
            Expected Improvement of X
        """
        if self.eta is None:
            raise ValueError('No current best specified. Call update('
                             'eta=<int>) to inform the acquisition function '
                             'about the current best value.')

        if len(X.shape) == 1:
            X = X[:, np.newaxis]

        m, var_ = self.model.predict_marginalized_over_instances(X)
        std = np.sqrt(var_)

        def calculate_log_ei():
            # we expect that f_min is in log-space
            f_min = self.eta - self.par
            v = (f_min - m) / std
            return (np.exp(f_min) * norm.cdf(v)) - \
                (np.exp(0.5 * var_ + m) * norm.cdf(v - std))

        if np.any(std == 0.0):
            # if std is zero, we have observed x on all instances
            # using a RF, std should be never exactly 0.0
            # Avoid zero division by setting all zeros in s to one.
            # Consider the corresponding results in f to be zero.
            self.logger.warning("Predicted std is 0.0 for at least one sample.")
            std_copy = np.copy(std)
            std[std_copy == 0.0] = 1.0
            log_ei = calculate_log_ei()
            log_ei[std_copy == 0.0] = 0.0
        else:
            log_ei = calculate_log_ei()

        if (log_ei < 0).any():
            raise ValueError(
                "Expected Improvement is smaller than 0 for at least one sample.")

        return log_ei.reshape((-1, 1))

class PI(AbstractAcquisitionFunction):
    def __init__(self,
                 model: AbstractEPM,
                 par: float = 0.0):
        r"""Computes the probability of improvement for a given x over the best so far value as acquisition value.

        :math:`P(f_{t+1}(\mathbf{X})\geq f(\mathbf{X^+}))` :math:`:= \Phi(\\frac{ \mu(\mathbf{X})-f(\mathbf{X^+}) }
        { \sigma(\mathbf{X}) })` with :math:`f(X^+)` as the incumbent and :math:`\Phi` the cdf of the standard normal

        Parameters
        ----------
        model : AbstractEPM
            A model that implements at least
                 - predict_marginalized_over_instances(X)
        par : float, default=0.0
            Controls the balance between exploration and exploitation of the
            acquisition function.
        """
        super(PI, self).__init__(model)
        self.long_name = 'Probability of Improvement'
        self.par = par
        self.eta = None
        self._required_updates = ('model', 'eta')

    def _compute(self, X: np.ndarray) -> np.ndarray:
        """Computes the PI value.

        Parameters
        ----------
        X: np.ndarray(N, D)
           Points to evaluate PI. N is the number of points and D the dimension for the points

        Returns
        -------
        np.ndarray(N,1)
            Expected Improvement of X
        """
        if self.eta is None:
            raise ValueError('No current best specified. Call update('
                             'eta=<float>) to inform the acquisition function '
                             'about the current best value.')

        if len(X.shape) == 1:
            X = X[:, np.newaxis]
        m, var_ = self.model.predict_marginalized_over_instances(X)
        std = np.sqrt(var_)
        return norm.cdf((self.eta - m - self.par) / std)

class LCB(AbstractAcquisitionFunction):
    def __init__(self,
                 model: AbstractEPM,
                 par: float = 1.0):
        r"""Computes the lower confidence bound for a given x over the best so far value as
        acquisition value.

        :math:`LCB(X) = \mu(\mathbf{X}) - \sqrt(\beta_t)\sigma(\mathbf{X})`

        Returns -LCB(X) as the acquisition_function optimizer maximizes the acquisition value.

        Parameters
        ----------
        model : AbstractEPM
            A model that implements at least
                 - predict_marginalized_over_instances(X)
        par : float, default=1.0
            Controls the balance between exploration and exploitation of the
            acquisition function.
        """
        super(LCB, self).__init__(model)
        self.long_name = 'Lower Confidence Bound'
        self.par = par
        self.num_data = None
        self._required_updates = ('model', 'num_data')

    def _compute(self, X: np.ndarray) -> np.ndarray:
        """Computes the LCB value.

        Parameters
        ----------
        X: np.ndarray(N, D)
           Points to evaluate LCB. N is the number of points and D the dimension for the points

        Returns
        -------
        np.ndarray(N,1)
            Expected Improvement of X
        """
        if self.num_data is None:
            raise ValueError('No current number of Datapoints specified. Call update('
                             'num_data=<int>) to inform the acquisition function '
                             'about the number of datapoints.')
        if len(X.shape) == 1:
            X = X[:, np.newaxis]
        m, var_ = self.model.predict_marginalized_over_instances(X)
        std = np.sqrt(var_)
        beta = 2 * np.log((X.shape[1] * self.num_data**2) / self.par)
        return -(m - np.sqrt(beta) * std)

class TS(AbstractAcquisitionFunction):
    def __init__(self,
                 model: AbstractEPM,
                 par: float = 0.0):
        r"""Do a Thompson Sampling for a given x over the best so far value as
        acquisition value.

        Thompson Sampling can only be used together with smac.optimizer.ei_optimization.RandomSearch, please do not
        use smac.optimizer.ei_optimization.LocalAndSortedRandomSearch to optimize TS acquisition function!!!

        :math:`TS(X) ~ \mathcal{N}(\mu(\mathbf{X}),\sigma(\mathbf{X}))'
        Returns -TS(X) as the acquisition_function optimizer maximizes the acquisition value.
        Parameters
        ----------
        model : AbstractEPM
            A model that implements at least
                 - predict_marginalized_over_instances(X)
        par : float, default=0.0
            TS does not require par here, we only wants to make it consistent with other acquisition functions
        """
        super(TS, self).__init__(model)
        self.long_name = 'Thompson Sampling'
        self.par = par
        self.num_data = None
        self._required_updates = ('model', )

    def _compute(self, X: np.ndarray) -> np.ndarray:
        """Sample a new value from a gaussian distribution whose mean and covariance values are given by model
        Parameters
        ----------
        X: np.ndarray(N, D)
           Points to be evaluated where we could sample a value. N is the number of points and D the dimension
           for the points
        Returns
        -------
        np.ndarray(N,1)
            negative sample value of X
        """
        if len(X.shape) == 1:
            X = X[:, np.newaxis]
        sample_function = getattr(self.model, "sample_functions", None)
        if callable(sample_function):
            return - sample_function(X, n_funcs=1)

        m, var_ = self.model.predict_marginalized_over_instances(X)
        rng = getattr(self.model, 'rng', np.random.RandomState(self.model.seed))
        m = m.flatten()
        var_ = np.diag(var_.flatten())
        return - rng.multivariate_normal(m, var_, 1).T


import logging
import os

import astropy.constants
import numpy as np
from astropy.coordinates import EarthLocation, SkyCoord
from astropy.io import fits
from astropy.io.ascii.cparser import CParserError  # pylint: disable=E0611
from astropy.table import Table, vstack
from astropy.time import Time
from easyquery import Query

from ..database import CsvTable, FitsTable
from .common import SPEED_OF_LIGHT, ensure_specs_dtype

__all__ = [
    "read_generic_spectra",
    "read_mmt",
    "read_mmt_bino",
    "read_aat",
    "read_aat_mz",
    "read_imacs",
    "read_wiyn",
    "read_palomar",
]

# pylint: disable=logging-format-interpolation

def get_obs_info_from_fits(
    fits_filepath,
    ra_name="RA",
    dec_name="DEC",
    time_name="MJD",
    ra_unit="hourangle",
    dec_unit="deg",
    time_format="mjd",
):
    hdr = fits.getheader(fits_filepath)
    sc = SkyCoord(hdr[ra_name], hdr[dec_name], unit=(ra_unit, dec_unit))
    obstime = Time(hdr[time_name], format=time_format)
    return sc, obstime

def heliocentric_correction(sc, obstime, site_name):
    helio_corr = sc.radial_velocity_correction(
        "heliocentric", obstime=obstime, location=EarthLocation.of_site(site_name)
    )
    return helio_corr.to_value(astropy.constants.c)  # pylint: disable=no-member

def read_generic_spectra(
    dir_path,
    extension,
    telname,
    usecols,
    n_cols_total,
    cuts=None,
    postprocess=None,
    fits_hdr_kwargs=None,
    helio_corr_site=None,
    before_time=None,
    table_read_kwargs=None,
    exclude_spec_masks=None,
    **kwargs,
):

    names = [usecols.get(i + 1, "_{}".format(i)) for i in range(n_cols_total)]
    exclude_names = [n for n in names if n.startswith("_")]

    output = []

    for filename in os.listdir(dir_path):
        filepath = os.path.join(dir_path, filename)
        rootname, extension_this = os.path.splitext(filename)

        if extension_this != extension:
            continue

        if "conflicted copy" in filename:
            logging.warning(
                "SKIPPING spectra file {} - it's a conflicted copy; check what went wrong!".format(filepath)
            )
            continue

        if exclude_spec_masks and (rootname in exclude_spec_masks or filename in exclude_spec_masks):
            continue

        helio_corr = None
        for fits_name in (
            rootname + ".fits.gz",
            rootname + ".fits",
            rootname.rpartition("_")[0] + ".fits.gz",
            rootname.rpartition("_")[0] + ".fits",
        ):
            fits_path = os.path.join(dir_path, fits_name)
            if os.path.isfile(fits_path):
                break
        try:
            sc, obstime = get_obs_info_from_fits(fits_path, **(fits_hdr_kwargs or {}))
        except (IOError, OSError):
            if fits_hdr_kwargs or helio_corr_site or before_time is not None:
                logging.warning("Cannot find or read corresponding fits file for {}".format(filepath))
        else:
            if before_time is not None and obstime > before_time:
                continue
            if helio_corr_site:
                helio_corr = heliocentric_correction(sc, obstime, helio_corr_site)

        try:
            table_read_kwargs_this = dict(
                format="ascii.fast_no_header",
                guess=False,
                names=names,
                exclude_names=exclude_names,
            )
            if table_read_kwargs:
                table_read_kwargs_this.update(table_read_kwargs)
            this = Table.read(filepath, **table_read_kwargs_this)
        except (IOError, CParserError) as e:
            logging.warning("SKIPPING spectra file {} - could not read or parse\n{}".format(filepath, e))
            continue

        this = ensure_specs_dtype(this, skip_missing_cols=True)
        this = Query(cuts).filter(this)
        if not len(this):
            continue

        if "MASKNAME" not in this.colnames:
            this["MASKNAME"] = filename

        if helio_corr is None:
            this["HELIO_CORR"] = False
        else:
            this["SPEC_Z"] += float(helio_corr)
            this["HELIO_CORR"] = True

        output.append(this)

    if not output:
        return

    output = vstack(output, "exact")
    output["TELNAME"] = telname
    if postprocess:
        output = postprocess(output)

    return ensure_specs_dtype(output)

def read_mmt(dir_path, before_time=None, exclude_spec_masks=None):
    extension = ".zlog"
    telname = "MMT"
    helio_corr_site = "mmt"

    n_cols_total = 11
    usecols = {
        2: "RA",
        3: "DEC",
        4: "mag",
        5: "SPEC_Z",
        6: "SPEC_Z_ERR",
        7: "ZQUALITY",
        8: "SPECOBJID",
    }

    cuts = Query("mag != 0", "ZQUALITY >= 0", (lambda x: x != "0", "SPECOBJID"))

    def postprocess(t):
        del t["mag"]
        t["RA"] *= 15.0
        return t

    return read_generic_spectra(**locals())

def read_mmt_bino(dir_path, exclude_spec_masks=None):
    extension = ".dat"
    telname = "BINO"

    n_cols_total = 8
    usecols = {
        1: "masknum",
        2: "SPECOBJID",
        3: "RA",
        4: "DEC",
        6: "SPEC_Z",
        7: "ZQUALITY",
    }

    table_read_kwargs = dict(
        format="ascii.fixed_width_no_header",
        col_starts=(0, 5, 11, 25, 38, 58, 66, 71),
    )

    def postprocess(t):
        t["SPEC_Z_ERR"] = 10 / SPEED_OF_LIGHT
        t["MASKNAME"] = np.char.add(np.char.add(t["MASKNAME"], "-"), t["masknum"].astype("<U"))
        del t["masknum"]
        return t

    return read_generic_spectra(**locals())

def read_aat(dir_path, before_time=None, exclude_spec_masks=None):
    extension = ".zlog"
    telname = "AAT"
    helio_corr_site = "sso"

    n_cols_total = 11
    usecols = {2: "RA", 3: "DEC", 5: "SPEC_Z", 7: "ZQUALITY", 8: "SPECOBJID"}
    cuts = Query("ZQUALITY >= 0", (lambda x: x != "0", "SPECOBJID"))

    fits_hdr_kwargs = dict(ra_name="MEANRA", dec_name="MEANDEC", time_name="UTMJD")

    def postprocess(t):
        t["SPEC_Z_ERR"] = 10 / SPEED_OF_LIGHT
        return t

    return read_generic_spectra(**locals())

def read_aat_mz(dir_path, before_time=None, exclude_spec_masks=None):
    extension = ".mz"
    telname = "AAT"
    helio_corr_site = "sso"

    n_cols_total = 15
    usecols = {3: "RA", 4: "DEC", 13: "SPEC_Z", 14: "ZQUALITY", 1: "SPECOBJID"}
    cuts = Query("ZQUALITY >= 0")

    fits_hdr_kwargs = dict(ra_name="MEANRA", dec_name="MEANDEC", time_name="UTMJD")
    table_read_kwargs = dict(delimiter=",")

    def postprocess(t):
        t["RA"] *= 180.0 / np.pi
        t["DEC"] *= 180.0 / np.pi
        t["SPEC_Z_ERR"] = 10 / SPEED_OF_LIGHT
        return t

    return read_generic_spectra(**locals())

def read_imacs(dir_path):
    extension = ".zlog"
    telname = "IMACS"

    n_cols_total = 12
    usecols = {
        2: "RA",
        3: "DEC",
        5: "SPEC_Z",
        6: "SPEC_Z_ERR",
        7: "ZQUALITY",
        8: "SPECOBJID",
        11: "MASKNAME",
    }

    cuts = Query("ZQUALITY >= 1", (lambda x: x != "0", "SPECOBJID"))

    return read_generic_spectra(**locals())

def read_wiyn(dir_path):

    output = []

    for f in os.listdir(dir_path):
        if not f.endswith(".fits.gz"):
            continue
        this = FitsTable(os.path.join(dir_path, f)).read()[["RA", "DEC", "ZQUALITY", "FID", "Z", "Z_ERR"]]
        this = Query("ZQUALITY >= 1").filter(this)
        this["MASKNAME"] = f
        output.append(this)

    output = vstack(output, "exact")

    output.rename_column("FID", "SPECOBJID")
    output.rename_column("Z", "SPEC_Z")
    output.rename_column("Z_ERR", "SPEC_Z_ERR")
    output["TELNAME"] = "WIYN"

    sc = SkyCoord(output["RA"], output["DEC"], unit=("hourangle", "deg"))
    output["RA"] = sc.ra.deg
    output["DEC"] = sc.dec.deg
    del sc

    return ensure_specs_dtype(output)

def read_palomar(file_path):
    if not hasattr(file_path, "read"):
        file_path = CsvTable(file_path)

    cols = [
        "SPECOBJID",
        "RA",
        "DEC",
        "SPEC_Z",
        "SPEC_Z_ERR",
        "ZQUALITY",
        "MASKNAME",
        "TELNAME",
        "HELIO_CORR",
    ]
    specs = file_path.read()[cols]

    return ensure_specs_dtype(specs)

import numpy,import astropy,from astropy