# The executable vikram_setup.py takes new onebody and twobody data
# comma-separate csv files and sets up the database aquarium.db
# This file uses it.

# command line with sqlite:
# sqlite3 aquarium.db
# To unlock an sqlite database type on the command line 
# mv aquarium.db tmp.db
# cp tmp.db aquarium.db

# To run with nohup, put 
# python3 vikram_protein.py > vikram.out 2> vikram.err
# in file runvik 
# chmod +x runvik
# nohup runvik &


# At position i=1 take the cost of that rotamer
# for position i=2 to i=n
#  for each rotamer k at position i
#     find the best cost up to k given the best cost up to position i-1 
#      combined with the best cost up to position i, rotamer k

# New approach looks not only at position i-1 but also at any position i'
# for i' < i that is next to i.
# So what we'd do is for all neighbors i' of i whose value is less than i,
#   for each rotamer j at position i
#     look at all combinations of positions of the neighbors i' 
#     and add up those costs

# But this might make the best rotamer for i' for i be different from the one
# for i' + 1.
# That is not good.    
# 
# Go through linearly but in the choice of rotamer at position i, also take 
# into account the pairwise cost with regard to all i' < i other than i-1.

# So, just go through linearly and then at each 
# position with more than two neighbors choose a new rotamer if that
# will reduce the total energy and then propagate that.
# To see if it does, just calculate the energy at that position 
# and then the pairwise energies with all neighbors. Call that the local delta

# We need a method that can rapidly calculate the total energy of a combination
# of positions.
# Once we get a profitable local delta, think about propagating that.

# We need a method that looks at other positions and computes the 
# best local delta.

# We can then do small perturbations. 
# Change a position that has many connections and see if that helps.
# So we can use simulated annealing
# If it does help then we can propagate.



import sys
import math
import csv
import os
import copy
import operator
import doctest
import collections
import datetime
import random
from operator import itemgetter, attrgetter
sys.setrecursionlimit(20000) # ??? might have to extend for different cities

import pandas as pd
import numpy as np
import sqlite3

now = datetime.datetime.now()
currentyear = now.year

#sqlite
print(datetime.datetime.now())
connection = sqlite3.connect("aquarium.db")
print(connection.total_changes)
cursor = connection.cursor()
rotnum = 0
rows = cursor.execute( "SELECT count(*) from singleenergy where rotamer = ? ", ([rotnum])).fetchall()
numpos = rows[0][0]
print("Number of positions is:")
print(numpos)

pos = 0
# rows = cursor.execute( "SELECT firstposition, firstrotamer, secondposition, secondrotamer, energy from crossenergy  where firstposition = ? and secondposition = firstposition+1", ([pos])).fetchall()
rows = cursor.execute( "SELECT count(*) from crossenergy").fetchall()
print("Number of rows in crossenergy: ")
print(rows)

# APPLICATION-SPECIFIC

# Perturb
# init = nodes that have more than two neighbors.
# Look at crossenergy and find all locations which have more than two neighbors.
# then we will add new ones if their neighbors have been modified
#
# For each such location i having rotamer r,  get all neighboring positions i' 
# For each such i' figure out its rotamer r'  as well as the cross
# costs with i for all rotamers or i with respect to i',r'
# Compute the local cost meaning cost of (i,r) + cost of (i,r) with each (i',r')
# Now imagine that you choose a different rotamer s for position i.
# Compute a new local cost. If its less than the current local cost 
# then choose that. Update the cost and the path.
#
def perturb():
    global bestpath
    global totalcost
    newones = [] # these will be added to when we change their neighbors
    print("faraways are: ")
    print(faraways) 
	# faraways are pairs in the format low position then high position
	# and in order
    neighbors = {}
    # for each position f[0] look for faraway neighbors greater than f[0]
    # for each position f[1] look for faraway neighbors less than f[1]
    print("Begin perturbing for faraways")
    print(datetime.datetime.now())
    for f in faraways:
      if f[0] in neighbors:
        neighbors[f[0]].append(f[1])
      else:
        neighbors[f[0]] = [f[1]]
      if f[1] in neighbors: 
        neighbors[f[1]].append(f[0])
      else:
        neighbors[f[1]] = [f[0]]
    # neighbors have a position p as a key  and then all faraway neighbors
    # of p
    print("neighbors are: ")
    print(neighbors)
      # neighbors is a map from positions to all faraway neighbors
    for loc in neighbors:
      # add in neighbors that are in sequence whether lower or higher
      if loc > 0:
        neighbors[loc].append(loc-1)
      if loc < len(bestpath) - 1:
        neighbors[loc].append(loc+1)
      # Simulated annealing might allow finding a better rotamer
      # even if the energy goes up. This might also give rise to several
      # new chromosomes, one for each rotamer at loc that looks promising.
      pair = findbetterrotamer(loc, neighbors[loc], bestpath)
      bestpath[loc] = pair[0] # better rotamer
      totalcost+= pair[1] # new cost
      if (0 > pair[1]):
        for n in neighbors[loc]:
          newones.append(n)
        # newones = list(dict.fromkeys(newones))
    # look at all the nodes that are neighbors of nodes that have 
    # have had their rotamers changed
    print("totalcost after adjusting the positions with multiple neighbors, but before relaxation: ")
    print(totalcost)
    print(datetime.datetime.now())
    relaxcounter = 0
    nochangecounter = 0
    flag = 1
    for pos in newones:
     if flag > 0:
      # newones = list(dict.fromkeys(newones))
      tmp = [x for x in newones if not x == pos]
      newones = copy.deepcopy(tmp)
      if pos not in neighbors:
        print("Trying position " + str(pos) + " which has no faroff neighbors.")
        if pos > 0:
          neighbors[pos].append(pos-1)
        if pos < len(bestpath) - 1:
          neighbors[pos].append(pos+1)
      else:
        print("Trying position " + str(pos) + " which has faroff neighbors.")
      pair = findbetterrotamer(pos, neighbors[pos], bestpath)
      bestpath[pos] = pair[0] # better rotamer
      totalcost+= pair[1] # new cost
      if (0 > pair[1]):
        nochangecounter = 0 # reset this to 0
        print("Relax counter: ")
        print(relaxcounter)
        print("Position: " + str(pos) + " New Rotamer: " + str(bestpath[pos]))
        print("Relax energy reduction: ")
        print(pair[1])
        relaxcounter+= 1
        print("Adding the following positions to newones: ")
        print(neighbors[pos])
        for n in neighbors[pos]:
          newones.append(n)
      else:
        print("Position " + str(pos) + " didn't have a better rotamer. ")
        nochangecounter+= 1
        if nochangecounter > nochangelimit:
          newones = []
          print("Too many times without progress")
          flag = 0 # stop this loop
    print("totalcost after relaxation: ")
    print("Length of newones is: " + str(len(newones)) + " and should be 0. ") 
    print("Final nochangecounter value was: " + str(nochangecounter))
    print(totalcost)
    
        

# Find a better rotamer at position than the one we have now
# This would be a good place for numpy
def findbetterrotamer(pos, neighs, bestpath):
  # fetch single cost at pos
  # then all cross costs with neighbors
  singlecost = cursor.execute( "SELECT rotamer, energy from singleenergy where position = ? ", ([pos])).fetchall()
  neighcost = {}
  # for pos, neighcost[neigh] will be all the pairwise costs between rotamers
  # at pos and the rotamer at neigh that is already chosen
  # note that neigh may have a lower value than pos
  for neigh in neighs:
    if (neigh > pos):
     rows = cursor.execute(
    "SELECT firstrotamer, secondrotamer, energy from crossenergy where firstposition = ? and secondposition = ? and secondrotamer = ?", [pos,neigh, bestpath[neigh]]).fetchall()
     neighcost[neigh] = rows
    elif (neigh < pos):
     rows = cursor.execute(
    "SELECT secondrotamer, firstrotamer, energy from crossenergy where firstposition = ? and secondposition = ? and firstrotamer = ?", [neigh,pos, bestpath[neigh]]).fetchall()
     neighcost[neigh] = rows
    # whether neigh < pos or neigh > pos the first field of rows will be the
    # one that can be changed, because that is the rotamer at pos
    # the second field will be the fixed position of the neighbor
  localcost = {}
  for spair in singlecost: 
    rot = spair[0]
    localcost[rot] = spair[1] # local cost for this rotamer at pos
  for n in neighcost:
    neightrip = neighcost[n]
    for t in neightrip:
      if t[0] in localcost:
        localcost[t[0]] += t[2]
      else:
        print("Neighbor has no crosscost with position " + str(t[0]))
  # find the minimum of the localcosts
  # change the rotamar here to that minimal rotamar
  # return the difference with the old cost and the new rotamer
  tmpnow = localcost[bestpath[pos]]
  min_val = min(localcost.values())
  L = list(localcost.items())
  for x in L:
    if x[1] == min_val:
      bestrot = x[0]
  return(bestrot, min_val - tmpnow)


    
    
    
# compute costs of all non-neighboring pairs
def  computecrosscost(bestpath): 
  global faraways  # for each i with faraway neighbors i' such that i' > i, find positions of those neighbors
  tempcrossout = []
  i = 0
  while (i < len(bestpath)):
    rows = cursor.execute(
    "SELECT firstposition, firstrotamer, secondposition, secondrotamer, energy  from crossenergy where firstposition = ? and firstrotamer = ? and secondposition - 1 > ? ", ([i,bestpath[i], i])).fetchall()
    paircostlist = rows
    if len(paircostlist) > 0: 
      tempfaraways = []
      tmpsum = 0
      for r in rows:
        if (r[3] == bestpath[r[2]]):
          faraways.append([r[0], r[2]])
          tmpsum+= r[4]
      tempcrossout.append(tmpsum)
    else:
      tempcrossout.append(0)
    i+=1
  return(tempcrossout)

# compute costs of all pairs, both neighboring and non-neighboring
def  computetotalcost(bestpath): 
  temptotal = []
  i = 0
  while (i < len(bestpath)):
    singlecost = cursor.execute( "SELECT energy from singleenergy where position = ? and rotamer = ?", ([i,bestpath[i]])).fetchall()
    tmpsum= singlecost[0][0]
    rows = cursor.execute(
    "SELECT firstposition, firstrotamer, secondposition, secondrotamer, energy  from crossenergy where firstposition = ? and firstrotamer = ? and secondposition  > ? ", ([i,bestpath[i], i])).fetchall()
    paircostlist = rows
    if len(paircostlist) > 0: 
      for r in rows:
        if (r[3] == bestpath[r[2]]):
          tmpsum+= r[4]
    temptotal.append(tmpsum)
    i+=1
  return(temptotal)


# 

# j is an index in the last rotamer that ends with the best cost
# We want to construct the set of all paths leading to j in the last location
# prevbest is the array which says which previous rotamer was best.
def constructpath(j):
  lastcosts = endcost[-1]
  allpathcosts = [lastcosts[j]]
  allpath = [j]
  loc = len(prevbest) - 1
  while (loc >= 0):
    prevj = prevbest[loc][allpath[-1]]
    allpath.append(prevj)
    allpathcosts.append(endcost[loc][prevj])
    loc-= 1
  return([allpath[::-1],allpathcosts[::-1]])
    

  

# dynamic programming algorithm that implements
# At position i=1 take the cost of that rotamer
# for position i=2 to i=n
#  for each rotamer k at position i
#     find the best cost up to k given the best cost up to position i-1
#      combined with the best cost up to position i, rotamer k
def bestbackbones():
  global endcost
  global prevbest
  temppos = 0
  rows = cursor.execute( "SELECT energy from singleenergy where position = ? ", ([temppos])).fetchall()
  y = [x[0] for x in rows]
  endcost.append(y)
  i = 1
  while (i < numpos ):
# starting at position 1 and going to the end
    temppos=i
    rows = cursor.execute( "SELECT energy from singleenergy where position = ? ", ([temppos])).fetchall()
    y = [x[0] for x in rows]
    mycostlist = copy.deepcopy(y) # individual costs of each rotamer at position i ignoring pairwise interactions
    prevcostlist = endcost[i-1] # costs up to previous rotamer
    prevbestlist = [-1] * len(mycostlist)  # will be best rotamer from position i-1 for each current rotamer at position i. This initializes to an impossible number
    target_firstposition = i-1
    target_secondposition = i
    rows = cursor.execute(
    "SELECT firstrotamer, secondrotamer, energy  from crossenergy where firstposition = ? and secondposition = ?", ([target_firstposition,target_secondposition])).fetchall()
    paircostlist = rows
    # costs of going from each k in position i-1 to each j in position i
    # Now we want to find the best path to each rotamer j of position i
    x1 = paircostlist
    for j in range(len(mycostlist)):
 # for each rotamer in position i
      pairstoj = []
      itest = 0
      for x2 in x1:
        if x2[1] == j:
          pairstoj.append([x2[0], x2[2]])
          if (x2[0] != itest):
            print(("Mismatch between rotamer index and array index:") + str(x2[0]) + "  "+ str(itest))
          itest+=1
      # breakpoint()
      mymin = mycostlist[j] + prevcostlist[0] + pairstoj[0][1]  # mycost +  the previous cost at position 0 + pairwise cost from that one at position 0
      myprevbest = 0  # initialize by saying previous best is 0
      for k in range(len(prevcostlist)): # for each k in previous position we want to go to position j 
        c = mycostlist[j] + prevcostlist[k] + pairstoj[k][1] # note that this works because each position k corresponds exactly to rotamer k (i.e. the k corresponds exactly to rotamer k (i.e. the mismatch above doesn't occur)
        if c < mymin:
          mymin = c
          myprevbest = k
      mycostlist[j] = mymin
      prevbestlist[j] = myprevbest
    # print("mycostlist for position " + str(i) )
    # print(mycostlist)
    endcost.append(mycostlist)
    prevbest.append(prevbestlist)
    i+=1

  




  

# DATA

poscost = [] # for position i, this gives the individual cost of each rotamer at that position

paircost = [] # for position i, this gives the pairwise cost of each rotamer at i with each rotamer at position i+1

endcost = [] # this is a list of lists. For the ith position , this gives the minimum cost to each rotamer at the ith position.


prevbest = [] # for the ith position and jth rotamer, it will give the best rotamer at position i-1 that gives the lowest cost to get to (i,j) 

# EXECUTION

bestbackbones() # find best rotamer choices looking only at the amino acids of the backbone


print(endcost)
print(prevbest)

lastcosts = endcost[-1]
overallmincost = min(lastcosts)
overallmincostindex = lastcosts.index(min(lastcosts))
pair = constructpath(overallmincostindex)
print("Best path from beginning to end considering only the backbone: " )
bestpath = pair[0]
print(bestpath)
print("Best path adjusted to start with 1 instead of zero" )
tempbestpath = [x+1 for x in bestpath]
print("Costs along that path, considering only the backbone " )
backbonecost = pair[1]
print(backbonecost)
faraways = [] # locations that have faraway neighbors
crosscosts = computecrosscost(bestpath) # compute cost to non-neighbor nodes
print("Totals of all cross-costs (i.e. costs that aren't neighbors on the backbone) " )
print(crosscosts)
# breakpoint()
totalcost = backbonecost[-1] + sum(crosscosts)
print("Total costs before perturbation (i.e. neighbors and non-neighbors) " )
print(totalcost)
# breakpoint()
nochangelimit = 2 * numpos # if there is no improvement after this number of attempts in a row, then give up
perturb()
print("Best path after perturbation " )
print(bestpath)
print("Best path adjusted to start with 1 instead of zero" )
tempbestpath = [x+1 for x in bestpath]
print(tempbestpath)
elements = computetotalcost(bestpath)
print("Position, rotamer, energy")
for i in range(len(bestpath)):
  print(str(i+1) + ": " +  str(tempbestpath[i]) + " " + str(elements[i]))

print("Total costs after perturbation " )
print(totalcost)
print("Here are the new crosscost costs total")
print(computecrosscost(bestpath))
print(datetime.datetime.now())
# print("Here are the total costs ")
# print(computetotalcost(bestpath))
# breakpoint()
connection.close()