(* This script calculates the Jensen-Shannon divergence (and the Jensen-Shannon metric) of the distribution form each other and from an uniform distributions. In order to calculate the JS integrals, the actin distributions are regularized (but NOT binned) on a finely-meshed grid of points. *)

(* Last update March 6 2018 *)

(* Set up string variables needed to read the distributions. 
NOTE: this may need to be adjusted by the user to account for different file naming conventions *)

(* neuronExperimentTag = "ABL KD"  ; timeTag ="3min" ; numberOfNeurons = 14 ; *)

(*  neuronExperimentTag = "ABL KD"  ; timeTag ="1min" ; numberOfNeurons = 13 ; *)

neuronExperimentTag = "ABL GOF" ;  timeTag ="3min" ; numberOfNeurons = 14 ;

(* neuronExperimentTag = "ABL GOF" ;  timeTag ="1min" ; numberOfNeurons = 13 ; *)

(* neuronExperimentTag = "WT" ; timeTag ="3min" ; numberOfNeurons = 14 ; *)

(* neuronExperimentTag = "WT" ; timeTag ="1min" ; numberOfNeurons = 14 ; *)

(* neuronExperimentTag = "Dicer" ; timeTag ="3min" ; numberOfNeurons = 14 ; *)

neuronNameList =  Table[   neuronExperimentTag  <> " " <> ToString[ nx ], { nx, 1, numberOfNeurons } ]

(* Set up some paths *)

pathToTraceDirectory =  "/Users/mcqueenp/ActiveProjects/AxonGrowthInDrosophilaWing/TraceData/" <> If[ neuronExperimentTag == "WT", "Wild Type",  neuronExperimentTag  ] <> " " <> timeTag <> "/"

pathToOutputDirectory =  "/Users/mcqueenp/Downloads/"

timeInfoFile = "IntervalBetweenFrames.csv"  ;

beginningPointOnAxonFile = "BeginningMicrometer.csv"   ;

actinFileSuffix= "_actin.csv"   ;

verySmallNumber = 10^-10

gridSeparation = 0.1

timeInfoFile = "IntervalBetweenFrames.csv"

beginningPointOnAxonFile = "BeginningMicrometer.csv"

actinFileName = "actin.csv"

neuronExperimentAll= {} ; timeIndexComparisonAll = {} ; jsMetricByNeuron = {} ; timeListByNeuron = {} ;jsMComparisonsAll = {} ;  timeIndexPlusOneAll = {} ; jsMTimePlusOneAll = {} ; neuronExperimentsPlusOneAll = {} ;entropyAll = {} ;filesWithFailedImports ={} ;

For[ nx = 1, nx ≤Length[ neuronNameList ], nx++,
pathToNeuronDirectory = pathToTraceDirectory <> neuronNameList[[ nx ]] <> "/" ;

(* Import time information data *)
fileName = pathToNeuronDirectory <> timeInfoFile ;
fileThere = FileExistsQ[ fileName ] ; 
If[ !fileThere, AppendTo[ filesWithFailedImports, fileName ]  ] ;
If[ !fileThere, Continue[ ] ] ;
dumArray = Import[ fileName ] ;
maximumTimeIndex = dumArray[[ 2 ]][[ 2 ]] ;
minimumTimeIndex = dumArray[[ 2 ]][[ 1 ]] ;
(* deltaTime = ( 60.0*dumArray[[ 2 ]][[ 3 ]] + dumArray[[ 2 ]][[ 1 ]])/60.0 ; *)
dumArray =. ;


beginningPointForAnalysis =0.0 ;
fileName = pathToNeuronDirectory <>  beginningPointOnAxonFile ;
fileThere = FileExistsQ[ fileName ] ;
If[ fileThere, dumArray = Import[ fileName ], AppendTo[ filesWithFailedImports, fileName ]   ] ;
If[ fileThere,beginningPointForAnalysis = dumArray[[ 2 ]][[ 1 ]] ] ;
If[ fileThere, dumArray =. ] ;



(* Import the raw actin data for each time point *)
timeIndexTX = {} ; actinDataRawTX = {} ; xCoordinateRawTX = {} ;
edgeLengthMax = -1000000.0 ;
For[ tx = minimumTimeIndex , tx ≤ maximumTimeIndex,  tx++,

timeString ="t=" <> ToString[ tx ] <> "_"  ;

fileName = pathToNeuronDirectory <> timeString <> actinFileName  ;
fileThere = FileExistsQ[ fileName ] ;
If[ !fileThere, AppendTo[ filesWithFailedImports, fileName ]  ] ;
If[ !fileThere, Continue[ ] ] ;

dumArray= Import[ fileName ] ;

actinData= SortBy[ Table[ { dumArray[[ n ]][[ 4 ]], dumArray[[ n ]][[ 5 ]] }, { n, 5, Length[ dumArray ] } ], First ] ;

AppendTo[ actinDataRawTX, actinData ] ;
AppendTo[ timeIndexTX, tx ] ;

locmax = Max[ Table[ actinData[[ ix ]][[ 1 ]], { ix, 1, Length[ actinData ] } ] ] ;
If[ edgeLengthMax < locmax, edgeLengthMax = locmax ] ;


dumArray =. ;
actinData =.;

] ; 

AppendTo[ timeListByNeuron , timeIndexTX ] ; 

numberOfGridPoints = Ceiling[(  edgeLengthMax - beginningPointForAnalysis )/gridSeparation ] + 1  ;


(* Regularize and normalize the intensities *)
actinIntensityRegularizedTX = {} ; entropyTX ={} ;
For[ tx = 1, tx ≤ Length[ timeIndexTX ], tx++, 

actIntReg = Table[ 0, { gg, 1, numberOfGridPoints } ] ;

For[ inx = 1, inx ≤ Length[  actinDataRawTX[[ tx ]] ] - 1, inx++, 

If[  actinDataRawTX[[ tx ]][[ inx ]][[ 1 ]] < beginningPointForAnalysis, Continue[ ] ]   ; 

gpA = Ceiling[ ( actinDataRawTX[[ tx ]][[ inx ]][[ 1 ]] - beginningPointForAnalysis )/gridSeparation ] + 1  ;
If[ gpA < 1, gpA = 1 ] ;
If[ gpA > numberOfGridPoints, gpA = numberOfGridPoints ] ;

gpB = Ceiling[(  actinDataRawTX[[ tx ]][[ inx + 1 ]][[ 1 ]]- - beginningPointForAnalysis  )/gridSeparation ] + 1  ;
If[ gpB< 1, gpB = 1 ] ;
If[ gpB > numberOfGridPoints, gpB = numberOfGridPoints ] ;


If[ gpA == gpB , actIntReg[[ gpA ]] += (  actinDataRawTX[[ tx ]][[ inx ]][[ 2 ]]  +  actinDataRawTX[[ tx ]][[ inx + 1 ]][[ 2 ]] ) ] ;
If[ gpA == gpB, Continue[ ] ] ;

ddA = 1.0*( actinDataRawTX[[ tx ]][[ inx + 1 ]][[ 2 ]] - actinDataRawTX[[ tx ]][[ inx  ]][[ 2 ]] )/( gpB - gpA ) ;
For[ gg = gpA, gg ≤ gpB, gg++, 
actIntReg[[ gg ]] += 1.0*actinDataRawTX[[ tx ]][[ inx  ]][[ 2 ]] + ( gg - gpA )*ddA ;
] ;

] ;



(* Use trapazoid approximation to integrate the intensities *)
actIntTrapAppx = 0 ;
For[ gg = 1, gg < numberOfGridPoints, gg++, 

deltaX =  gridSeparation ; 
actLocal =  actIntReg[[ gg + 1  ]] + actIntReg[[ gg   ]] ;
actIntTrapAppx  += 0.5*deltaX*actLocal ;

] ;

For[ gg = 1, gg <= numberOfGridPoints, gg++, actIntReg[[ gg ]] = actIntReg[[ gg ]]/actIntTrapAppx ] ;

AppendTo[ actinIntensityRegularizedTX, actIntReg ] ;



(* Find entropies of the regularized, normalized distributions, using trapezoid approximation *)
entropy = 0.0 ;
For[ gg = 1, gg < numberOfGridPoints, gg++, 

deltaX =  gridSeparation ; 
actLocal =  If[ actIntReg[[ gg + 1  ]] ≤ 0.0, 0.0,  actIntReg[[ gg + 1  ]] *Log[ 2.0,  actIntReg[[ gg + 1  ]]  ] ] 
+ If[ actIntReg[[ gg   ]] ≤ 0.0, 0.0,  actIntReg[[ gg   ]]*Log[ 2.0, actIntReg[[ gg   ]] ] ] ;
actLocal *= -1.0 ;

entropy +=  0.5*deltaX*actLocal  ;

] ; 

AppendTo[ entropyTX, entropy ] ;

actIntReg =.; 

] ;

AppendTo[ entropyAll, entropyTX ] ;


jsMetricByNeuronLocal = Table[ Table[ 0,{ txI, 1,  Length[ timeIndexTX ] } ], { txO, 1,  Length[ timeIndexTX ] } ] ;

For[ txO = 1, txO < Length[ timeIndexTX ], txO++, 
For[ txI = txO + 1, txI ≤ Length[ timeIndexTX ], txI++,

jsD = 0.0 ;
For[ gg = 1, gg < numberOfGridPoints, gg++,

pppA = 0.5*(  actinIntensityRegularizedTX[[ txI  ]][[ gg ]] +  actinIntensityRegularizedTX[[ txO ]][[ gg ]] )  ;
pppB= 0.5*(  actinIntensityRegularizedTX[[ txI  ]][[ gg  + 1 ]] + actinIntensityRegularizedTX[[ txO ]][[ gg  + 1 ]] )  ;

pppLoc = If[ pppA ≤ 0.0, 0.0, pppA*Log[ 2.0, pppA ] ] + If[ pppB ≤ 0.0, 0.0, pppB*Log[ 2.0, pppB ] ] ;
pppLoc *= -1.0 ;

deltaX =  gridSeparation ; 
jsD +=  0.5*deltaX*pppLoc  ;

] ;

jsD  += -0.5*( entropyTX[[ txO ]] + entropyTX[[ txI ]] ) ;
jsM = If[ jsD > 0.0, jsD^0.5, 0.0 ] ;

AppendTo[ neuronExperimentAll, neuronNameList[[ nx ]] ] ;
AppendTo[ timeIndexComparisonAll, ToString[ timeIndexTX[[ txO ]] ] <> " - " <> ToString[ timeIndexTX[[ txI ]] ] ] ;
AppendTo[ jsMComparisonsAll, jsM ] ;

jsMetricByNeuronLocal[[ txI ]][[ txO ]] =  jsM ;
jsMetricByNeuronLocal[[ txO ]][[ txI ]] =  jsM ;


If[  timeIndexTX[[ txO ]] ≠ timeIndexTX[[ txI ]] - 1, Continue[ ] ] ;

AppendTo[ timeIndexPlusOneAll , ToString[ timeIndexTX[[ txO ]] ] <> " - "<> ToString[ timeIndexTX[[ txI ]] ] ] ;
AppendTo[ jsMTimePlusOneAll , jsM ] ;
AppendTo[  neuronExperimentsPlusOneAll, neuronNameList[[ nx ]] ] ;


] ;

] ;


AppendTo[ jsMetricByNeuron,    jsMetricByNeuronLocal ] ; 

entropyTX =.;
actinIntensityRegularizedTX =.; 

] 
(* End of loop on nx *)

filesWithFailedImports

ListPlot[ Flatten[  entropyAll ]  ]

ListPlot[ jsMTimePlusOneAll ]

ListPlot[ Transpose[{ Sort[  jsMTimePlusOneAll ], Table[ ix/Length[ jsMTimePlusOneAll ], { ix, 1, Length[ jsMTimePlusOneAll ] } ] }] ]

ListPlot[  jsMComparisonsAll ]

{ Length[ entropyAll ], Length[ Flatten[ entropyAll ] ], Length[ timeIndexPlusOneAll ], Length[ jsMTimePlusOneAll ], Length[ neuronExperimentsPlusOneAll ], Length[ neuronExperimentAll ], Length[ timeIndexComparisonAll ], Length[ jsMComparisonsAll ] }

dumArray = {{ "Neuron", "t - t+1",  "Comparison t to t+1 JS Divergence", "Comparison t to t+1 JS Metric" } } ;

For[ inx = 1, inx ≤ Length[  neuronExperimentsPlusOneAll ], inx++,

  jsD = jsMTimePlusOneAll[[ inx ]]^2 ;

jsM =  jsMTimePlusOneAll[[ inx ]] ;

lineArray = { neuronExperimentsPlusOneAll[[ inx ]], timeIndexPlusOneAll[[ inx ]],  jsD, jsM } ;

AppendTo[ dumArray, lineArray ] ; 

]

Length[ dumArray ]

dumArray[[ 6 ]]

ListPlot[ dumArray[[ All, 3 ]] ]

ListPlot[ dumArray[[ All, 4 ]] ]

Export[  pathToOutputDirectory <> neuronExperimentTag <> "_" <> timeTag <> "_JSMetricTTPlusOne.csv", dumArray, "CSV" ]

{  Length[ neuronExperimentAll ], Length[ timeIndexComparisonAll ], Length[ jsMComparisonsAll ] }

dumArray = {{ "Neuron", "Times Compared",  "JS Divergence", "JS Metric" } } ;

For[ inx = 1, inx ≤ Length[  neuronExperimentAll ], inx++,

jsD = jsMComparisonsAll[[ inx ]]^2  ;
jsM = jsMComparisonsAll[[ inx ]] ;

lineArray = { neuronExperimentAll[[ inx ]], timeIndexComparisonAll[[ inx ]],  jsD, jsM } ;

AppendTo[ dumArray, lineArray ] ; 

]

Length[ dumArray ]

dumArray[[ 6 ]]

ListPlot[ dumArray[[ All, 3 ]] ]

ListPlot[ dumArray[[ All, 4 ]] ]

Export[  pathToOutputDirectory <> neuronExperimentTag <> "_" <> timeTag <> "_JSMetric_All_Comparisons.csv", dumArray, "CSV" ]

(* Make color scale for the JS metric displays *)

colorJS[ x_ ] := If[ x ≤ 0, Black, 
If[ x ≤ 1/10, RGBColor[ 0, 0.25, 0.25 ],
If[ x ≤ 2/10, RGBColor[ 0, 0.5, 0.5 ],
If[ x ≤ 3/10, RGBColor[ 0, 0.75, 0.25  ],
If[ x ≤  4/10, RGBColor[ 1, 1, 0.5  ],


If[ x ≤ 1/2, RGBColor[  1, 0.85, 0.45  ],


If[ x ≤ 6/10, RGBColor[ 0.75, 0.55, 0 ],
If[ x ≤ 7/10, RGBColor[  1, 0.5, 0 ],
If[ x ≤ 8/10, RGBColor[ 0.75, 0, 0 ],
If[ x ≤ 9/10, RGBColor[ 0.5, 0.125, 0 ],
If[ x ≤ 1, RGBColor[ 0.5, 0, 0.5 ],
RGBColor[ 1, 1, 1 ] 
 ]
]
]
]
]
]
]
]
]
]
]

colorArray = Table[ Table[ n*0.01, { ix, 0, 11 } ], { n, 0, 101 } ] ;

arrPl = ArrayPlot[ colorArray , ColorFunction -> (colorJS[#]&), ColorFunctionScaling->False, AspectRatio -> 10, Frame -> True, FrameStyle->Directive[ Opacity[0],FontSize[ 14 ], Thick ],   FrameTicksStyle -> Directive[ Opacity[1], 30   ],  FrameTicks ->  { { None, Table[ { n + 0.5, ToString[ 1.0*( n - 1 )/100 ]  }, { n, 1, 101, 10 } ] } ,  { None, None } },
 ImageSize -> 150   ]

Export[ pathToOutputDirectory <> "JSComparison_ColorCode.tif", arrPl, "Tiff" ]

For[ nx = 1, nx ≤ Length[ neuronNameList  ], nx++,

mpl =  ArrayPlot[  jsMetricByNeuron[[ nx ]],  Frame ->  True, ColorFunction-> (colorJS[#]&), ColorFunctionScaling->False, ImageSize -> ( Length[  timeListByNeuron [[ nx ]] ] + 10 )*400/14, PlotRange -> { { -10, Length[  timeListByNeuron[[ nx ]] ] }, { 0, Length[ timeListByNeuron[[ nx ]]  ] + 10 } }, 
BaseStyle->{ 15, FontFamily->"Arial"},
FrameStyle->Directive[ Opacity[0],FontSize[ 14 ], Thick ],   FrameTicksStyle -> Directive[ Opacity[1], 30   ],  FrameTicks ->  { {  Table[ { tx , ToString[ tx - 1 ] },{ tx, 1, Length[ timeListByNeuron[[ nx  ]] ], 5 } ],  None   }, { None, Table[ { tx , ToString[ tx - 1 ] },{ tx, 1, Length[ timeListByNeuron[[ nx ]] ], 5 } ] } }

] ;


mplD =  ArrayPlot[  jsMetricByNeuron[[ nx ]]^2,  Frame ->  True, ColorFunction-> (colorJS[#]&), ColorFunctionScaling->False, ImageSize -> ( Length[  timeListByNeuron [[ nx ]] ] + 10 )*400/14, PlotRange -> { { -10, Length[  timeListByNeuron[[ nx ]] ] }, { 0, Length[ timeListByNeuron[[ nx ]]  ] + 10 } }, 
BaseStyle->{ 15, FontFamily->"Arial"},
FrameStyle->Directive[ Opacity[0],FontSize[ 14 ], Thick ],   FrameTicksStyle -> Directive[ Opacity[1], 30   ],  FrameTicks ->  { {  Table[ { tx , ToString[ tx - 1 ] },{ tx, 1, Length[ timeListByNeuron[[ nx  ]] ], 5 } ],  None   }, { None, Table[ { tx , ToString[ tx - 1 ] },{ tx, 1, Length[ timeListByNeuron[[ nx ]] ], 5 } ] } }

] ;

Export[ pathToOutputDirectory <> "JSMetric_" <>  neuronNameList[[ nx ]] <>  ".tif"  , mpl, "Tiff" ] ;
Export[ pathToOutputDirectory <> "JSDivergence_" <>  neuronNameList[[ nx ]] <>  ".tif"  , mplD, "Tiff" ] ;

]

(* End of file *)