"""
A script to estimate the HIV epidemic model parameters using ABC.
"""

from apgl.util import *
from exp.viroscopy.model.HIVGraph import HIVGraph
from exp.viroscopy.model.HIVABCParameters import HIVABCParameters
from exp.viroscopy.model.HIVEpidemicModel import HIVEpidemicModel
from exp.viroscopy.model.HIVRates import HIVRates
from exp.viroscopy.model.HIVModelUtils import HIVModelUtils
from exp.viroscopy.model.HIVGraphMetrics import HIVGraphMetrics2
from apgl.predictors.ABCSMC import ABCSMC

import logging
import sys
import numpy
import multiprocessing

FORMAT = "%(levelname)s:root:%(process)d:%(message)s"
logging.basicConfig(stream=sys.stdout, level=logging.DEBUG, format=FORMAT)
numpy.set_printoptions(suppress=True, precision=4, linewidth=100)
numpy.seterr(invalid='raise')

#First try the experiment on some toy data 
resultsDir = PathDefaults.getOutputDir() + "viroscopy/toy/" 
graphFile = resultsDir + "ToyEpidemicGraph0"
targetGraph = HIVGraph.load(graphFile)

numTimeSteps = 10 
T, recordStep, printStep, M = HIVModelUtils.defaultSimulationParams()
times = numpy.linspace(0, T, numTimeSteps)
graphMetrics = HIVGraphMetrics2(times)

realSummary = graphMetrics.summary(targetGraph)
epsilonArray = numpy.array([0.8, 0.6, 0.5])*numTimeSteps

def breakFunc(graph, currentTime): 
    return graphMetrics.shouldBreak(realSummary, graph, epsilonArray[0], currentTime)

def createModel(t):
    """
    The parameter t is the particle index. 
    """
    undirected = True
    T, recordStep, printStep, M = HIVModelUtils.defaultSimulationParams()
    graph = HIVGraph(M, undirected)
    
    alpha = 2
    zeroVal = 0.9
    p = Util.powerLawProbs(alpha, zeroVal)
    hiddenDegSeq = Util.randomChoice(p, graph.getNumVertices())

    rates = HIVRates(graph, hiddenDegSeq)
    model = HIVEpidemicModel(graph, rates, T)
    model.setRecordStep(recordStep)
    model.setPrintStep(printStep)
    model.setBreakFunction(breakFunc) 

    return model

if len(sys.argv) > 1:
    numProcesses = int(sys.argv[1])
else: 
    numProcesses = multiprocessing.cpu_count()

posteriorSampleSize = 20
thetaLen = 10

logging.debug("Posterior sample size " + str(posteriorSampleSize))

meanTheta = HIVModelUtils.defaultTheta()
abcParams = HIVABCParameters(meanTheta, 0.5, 0.2)

abcSMC = ABCSMC(epsilonArray, realSummary, createModel, abcParams, graphMetrics)
abcSMC.setPosteriorSampleSize(posteriorSampleSize)
thetasArray = abcSMC.run()

meanTheta = numpy.mean(thetasArray, 0)
stdTheta = numpy.std(thetasArray, 0)
logging.debug(thetasArray)
logging.debug("meanTheta=" + str(meanTheta))
logging.debug("stdTheta=" + str(stdTheta))
logging.debug("realTheta=" + str(HIVModelUtils.defaultTheta()))

thetaFileName =  resultsDir + "ThetaDistSimulated.pkl"
Util.savePickle(thetasArray, thetaFileName)

@@ -424,16 +424,17 @@
 hMetrics
+2
  import 
@@ -450,16 +450,118 @@
 etrics2%0A
+from exp.viroscopy.model.HIVVertices import HIVVertices%0Afrom exp.sandbox.GraphMatch import GraphMatch%0A
 from apg
@@ -658,16 +658,55 @@
 essing%0A%0A
+assert False, %22Must run with -O flag%22%0A%0A
 FORMAT =
@@ -927,52 +927,8 @@
 ')%0A%0A
-#First try the experiment on some toy data %0A
 resu
@@ -1158,16 +1158,43 @@
 arams()%0A
+startDate = 0%0AendDate = T%0A%0A
 times = 
@@ -1232,96 +1232,8 @@
 ps)%0A
-graphMetrics = HIVGraphMetrics2(times)%0A%0ArealSummary = graphMetrics.summary(targetGraph)%0A
 epsi
@@ -1260,21 +1260,16 @@
 ay(%5B
-0.8, 
 0.6, 0.5
 %5D)*n
@@ -1268,145 +1268,14 @@
  0.5
-%5D)*numTimeSteps%0A%0Adef breakFunc(graph, currentTime): %0A    return graphMetrics.shouldBreak(realSummary, graph, epsilonArray%5B0%5D, currentTime
+, 0.3%5D
 )%0A%0Ad
@@ -1379,82 +1379,8 @@
 rue%0A
-    T, recordStep, printStep, M = HIVModelUtils.defaultSimulationParams()%0A
     
@@ -1554,16 +1554,400 @@
 tices())
+%0A    %0A    featureInds= numpy.ones(graph.vlist.getNumFeatures(), numpy.bool)%0A    featureInds%5BHIVVertices.infectionTimeIndex%5D = False %0A    featureInds%5BHIVVertices.hiddenDegreeIndex%5D = False %0A    featureInds = numpy.arange(featureInds.shape%5B0%5D)%5BfeatureInds%5D%0A    matcher = GraphMatch(%22U%22, featureInds=featureInds)%0A    graphMetrics = HIVGraphMetrics2(targetGraph, epsilonArray%5Bt%5D, matcher)
 %0A%0A    ra
@@ -2026,16 +2026,74 @@
 rates, T
+=float(endDate), T0=float(startDate), metrics=graphMetrics
 )%0A    mo
@@ -2159,47 +2159,8 @@
 tep)
-%0A    model.setBreakFunction(breakFunc) 
 %0A%0A  
@@ -2521,21 +2521,8 @@
 ray,
- realSummary,
  cre
@@ -2544,22 +2544,8 @@
 rams
-, graphMetrics
 )%0Aab

Use relevant feature inds, and update for new style ABC and metrics.