## ---------------------------------------------------------------------------
## MOI.py
## Last modified on: 2014-07-11
## M. Cracknell
## Polyspectral Services
## polyspectral.services@gmail.com
## ---------------------------------------------------------------------------

## modules
import os
from math import *
import numpy as np
from numpy import *
import scipy as sp
from scipy import linalg
import datetime
import csv

## must have vector.py and associated module vector.pyc installed in working directory 
import vector

## arcpy
import arcpy
import arcpy.mapping
## requires spatial analyst license
from arcpy.sa import *

## matplotlib
import matplotlib.pyplot as plt
import matplotlib.path as mpath
import matplotlib.patches as mpatch

## allow for overwrite
arcpy.env.overwriteOutput = True
## Check out any necessary licenses
arcpy.CheckOutExtension("spatial")

## ----------------------------------

## functions to run PointsToPlane

## to hold x, y and z coordinates
class Coord:
        def __init__(self, x, y, z):
                self.Xm = round(x,1)
                self.Ym = round(y,1)
                self.Zm = round(z,1)

## diffenence between mean and 3D points
class Diff:
        def __init__(self, x, y, z):
                self.Xdiff = x
                self.Ydiff = y
                self.Zdiff = z


## output data
class OutData:
        def __init__(self):
                self.ID = None
                self.N_points = None
                self.Xmean = None
                self.Ymean = None
                self.Zmean = None
                self.M = None
                self.K = None
                self.A = None
                self.B = None
                self.C = None
                self.D = None
                self.RMSE = None
                self.Dip = None
                self.Dipdir = None


## to hold 3D point coords for a single feature 
class Record:
        def __init__(self, newID):
                self.recID = newID
                self.coords = []
                self.diffs = []
                self.output = OutData()
                

        ## calculate average (x,y, and z) of input 3D points
        def average(self):
                xvals = 0.0
                yvals = 0.0
                zvals = 0.0
                count = 0
                for c in self.coords:
                        xvals = c.Xm + xvals
                        yvals = c.Ym + yvals
                        zvals = c.Zm + zvals
                        count = count + 1

                if count > 2:        
                        self.output.Xmean = xvals / count
                        self.output.Ymean = yvals / count
                        self.output.Zmean = zvals / count
                        self.output.N_points = count
                        self.output.ID = self.recID
                        
                        for c in self.coords:
                                xvals = c.Xm - self.output.Xmean
                                self.Xdiff = xvals
                                yvals = c.Ym - self.output.Ymean
                                self.Ydiff = yvals
                                zvals = c.Zm - self.output.Zmean
                                self.Zdiff = zvals
                                self.diffs.append(Diff(self.Xdiff,self.Ydiff,self.Zdiff))
                else:
                        self.output.Xmean = -9999
                        self.output.Ymean = -9999
                        self.output.Zmean = -9999
                        self.output.N_points = count
                        self.output.ID = self.recID

                        for c in self.coords:
                                self.diffs.append(Diff(-9999,-9999,-9999))

        ## MOI plane estimation
        ## from Fernandez (2005)
        def plane_calc(self):
                v11 = []
                v12 = []
                v13 = []
                v22 = []
                v23 = []
                v33 = []
                for p in self.diffs:
                        a = p.Xdiff
                        b = p.Ydiff
                        c = p.Zdiff
                        temp11 = a**2
                        v11.append(temp11)
                        temp12 = a * b
                        v12.append(temp12)
                        temp13 = a * c
                        v13.append(temp13)
                        temp22 = b**2
                        v22.append(temp22)
                        temp23 = b * c
                        v23.append(temp23)
                        temp33 = c**2
                        v33.append(temp33)

                sum11 = sum(v11)
                sum12 = sum(v12)
                sum13 = sum(v13)
                sum22 = sum(v22)
                sum23 = sum(v23)
                sum33 = sum(v33)
                matrix1 = [sum11, sum12, sum13]
                matrix2 =  [sum12, sum22, sum23]
                matrix3 = [sum13, sum23, sum33]
                la,v = linalg.eig([matrix1, matrix2, matrix3])
                l1, l2, l3 = la
                Eigenvalues = []
                Eigenvalues.append(l1.real)
                Eigenvalues.append(l2.real)
                Eigenvalues.append(l3.real)
                eigen3 = np.amin(Eigenvalues)
                posi = Eigenvalues.index(eigen3)
                v1, v2, v3 = v
                vec = [v1[posi],v2[posi],v3[posi]]
                Eigval = []
                if l1.real < l2.real and l2.real < l3.real:
                        eig1 = l3.real
                        eig2 = l2.real
                        eig3 = l1.real
                        Eigval = [eig1,eig2,eig3]
                elif l1.real < l2.real and l2.real > l3.real and l1.real > l3.real:
                        eig1 = l2.real
                        eig2 = l1.real
                        eig3 = l3.real
                        Eigval = [eig1,eig2,eig3]
                elif l1.real < l2.real and l2.real > l3.real and l1.real < l3.real:
                        eig1 = l2.real
                        eig2 = l3.real
                        eig3 = l1.real
                        Eigval = [eig1,eig2,eig3]
                elif l1.real > l2.real and l2.real < l3.real and l1.real < l3.real:
                        eig1 = l3.real
                        eig2 = l1.real
                        eig3 = l2.real
                        Eigval = [eig1,eig2,eig3]
                elif l1.real > l2.real and l2.real < l3.real and l1.real > l3.real:
                        eig1 = l1.real                                                                                                                                                                   
                        eig2 = l3.real
                        eig3 = l2.real
                        Eigval = [eig1,eig2,eig3]
                else:
                        eig1 = l1.real
                        eig2 = l2.real
                        eig3 = l3.real
                        Eigval = [eig1,eig2,eig3]
                        
                self.output.M = np.log(Eigval[0] / Eigval[2])
                self.output.K = np.log(Eigval[0] / Eigval[1]) / np.log(Eigval[1] / Eigval[2])  
                A, B, C, D = vector.plane([self.output.Xmean, self.output.Ymean, self.output.Zmean], vec)

                self.output.A = A
                self.output.B = B
                self.output.C = C
                self.output.D = D
                self.output.Eig = Eigval


        ## calculate geometric Root Mean Squared Error
        ## distance (m) from MOI plane to 3D points 
        ## from Stewart, 1999 p. 851 (find the perpendicular distance
        def RMSE(self):
                RMS = []
                for p in self.coords:                
                        num = (self.output.A * p.Xm) + (self.output.B * p.Ym) + (self.output.C * p.Zm) + self.output.D
                        den = sqrt((self.output.A ** 2) + (self.output.B ** 2) + (self.output.C ** 2))
                        distance = num / den 
                        RMS.append(distance**2)

                RMSError = sqrt(sum(RMS) / len(RMS))
                self.output.RMSE = RMSError
                self.output.ptRMSE = sqrt(RMS)

        ## calculate plane geometric orientation
        ## dip and dip direction
        ## equations from Vacher (2000)
        def orientation(self):
                ## strike orientation (Eq. #26)
                if self.output.A == 0 or self.output.B == 0 or self.output.C == 0:
                        self.output.Dip = -9999
                        self.output.Dipdir = -9999
                else:   
                        ## dip section (Eq. #27)
                        dipdir = atan(-self.output.A/self.output.B)
                        dipdir = round(dipdir * 180 / np.pi,0)
                      
                        ## dip angle (Eq. #30)
                        dip = atan(sqrt((self.output.A**2 + self.output.B**2)/self.output.C**2))
                        dip = round(dip * 180 / np.pi,0)
                        self.output.Dip = dip

                        ## set up partial derivatives to establish dip direction
                        ## (commments on p. 33)
                        dzdx = -self.output.A/self.output.C
                        dzdy = -self.output.B/self.output.C
                        if dzdx < -0.01 and dzdy > 0:
                                dipdir = 180 - abs(dipdir)
                                self.output.Dipdir = dipdir
                        elif dzdx > 0 and dzdy > 0:
                                dipdir = 180 + abs(dipdir)
                                self.output.Dipdir = dipdir
                        elif dzdx > 0 and dzdy < 0:
                                dipdir = 360 - abs(dipdir)
                                self.output.Dipdir = dipdir
                        elif dzdx < 0 and dzdy < 0:
                                if dipdir < 0:
                                        dipdir = 0 + abs(dipdir)
                                else:
                                        dipdir = 90 - dipdir
                                self.output.Dipdir = dipdir
                        elif dzdx < 0.01 and dzdx > -0.01 and dzdy > 0:
                                dipdir = 180
                        elif dzdx == 0 and dzdy < 0:
                                dipdir = 0
                        elif dzdx > 0 and dzdy == 0:
                                dipdir = 270
                        else: #dzdx < 0 and dzdy == 0:
                                dipdir = 90
                        self.output.Dipdir = dipdir


        ## write to output file
        def writer(self,outFC):
                cur = arcpy.UpdateCursor(outFC)
                ID = self.output.ID
                for row in cur:
                        if row.getValue('OBJECTID') == ID:
                            row.setValue('ID', ID)
                            row.setValue('N_points',self.output.N_points)
                            row.setValue('Xmean', round(self.output.Xmean,1))
                            row.setValue('Ymean', round(self.output.Ymean,1))
                            row.setValue('Zmean', round(self.output.Zmean,1))
                            row.setValue('M', round(self.output.M,3))
                            row.setValue('K', round(self.output.K,3))
                            row.setValue('RMSE', round(self.output.RMSE,1))
                            row.setValue('A', round(self.output.A,3))
                            row.setValue('B', round(self.output.B,3))
                            row.setValue('C', round(self.output.C,3))
                            row.setValue('D', round(self.output.D,0))
                            row.setValue('Dip', self.output.Dip)
                            row.setValue('Dipdir', self.output.Dipdir)
                            cur.updateRow(row)
                            ## outputs
                            return(ID,self.output.N_points,
                                round(self.output.Xmean,1),
                                round(self.output.Ymean,1),
                                round(self.output.Zmean,1),
                                round(self.output.M,3),
                                round(self.output.K,3),
                                round(self.output.RMSE,1),
                                round(self.output.A,3),
                                round(self.output.B,3),
                                round(self.output.C,3),
                                round(self.output.D,3),
                                self.output.Dip,
                                self.output.Dipdir,
                                self.output.Eig[0],
                                self.output.Eig[1],
                                self.output.Eig[2])
                                   
        ## output individualk point error for analysis                
        def Error(self):
                return self.output.ptRMSE

arcpy.AddMessage("Libraries and functions loaded")

## ----------------------------------

## Script arguments

## set up scratch and work geodatabase directorys
## scratch 
MOI_scratch_gdb = arcpy.GetParameterAsText(0)
if MOI_scratch_gdb == '#' or not MOI_scratch_gdb:
        MOI_scratch_gdb = "C:\ArcGIS\MOI\MOI_scratch.gdb" # provide a default value if unspecified
## work
MOI_template_gdb = arcpy.GetParameterAsText(1)
if MOI_template_gdb == '#' or not MOI_template_gdb:
        MOI_template_gdb = "C:\ArcGIS\MOI\MOI_template.gdb" # provide a default value if unspecified

## settings for manipulating map layers
## get the map document 
mxd = arcpy.mapping.MapDocument("CURRENT")  
## get the data frame 
df = arcpy.mapping.ListDataFrames(mxd,"*")[0]

## remove layers to be creasted again
for lyr in arcpy.mapping.ListLayers(mxd, "*",df):
        if lyr.name == "GeoPoint" or lyr.name == "points3D":
                arcpy.mapping.RemoveLayer(df, lyr)

arcpy.RefreshActiveView()
arcpy.RefreshTOC() 

        
## get vertices of GeoLines and create points
## then calc midpoint of line for plotting 
## based on code written by C. Mazengarb

arcpy.AddMessage("Loading 3D points ... ")

## input line feature
lineFC = "C:\\ArcGIS\\MOI\\MOI_template.gdb\\GeoLines"
midPoints = []
allVert = []

## get point coordinates 
dsc = arcpy.Describe(lineFC)
shapeFieldName = dsc.ShapeFieldName
updateLine = arcpy.UpdateCursor(lineFC)

#cur = arcpy.UpdateCursor(pointFC)

## loop through lines
count = 1
for row in updateLine:
        vertexList = []
        ## set line ID
        FID = row.getValue("OBJECTID")
        updateLine.updateRow(row)
        shapeObj = row.getValue(shapeFieldName)
        partObj = shapeObj.getPart(0)
        ## get coords
        for pointObj in partObj:
                vertexList.append((pointObj.X, pointObj.Y))
                allVert.append((count,FID,pointObj.X, pointObj.Y))
                count = count + 1
                
        vertexCount = len(vertexList)
        ## get midpoints
        LineMidpoint = int(round(vertexCount/2))
        CentreX = vertexList[LineMidpoint][0]
        CentreY = vertexList[LineMidpoint][1]
        midPoints.append((FID, CentreX, CentreY))

## output 3Dpoint features
## delete exisiting file
try:
        arcpy.Delete_management("C:\\ArcGIS\\MOI\\MOI_scratch.gdb\\temp3D")
except:
        print "feature class does not exist"

pt3D = arcpy.CreateFeatureclass_management("C:\\ArcGIS\\MOI\\MOI_scratch.gdb\\", "temp3D", "POINT", has_m="DISABLED", has_z="ENABLED", spatial_reference=lineFC)

## create new points
point = arcpy.Point()
ptGeom = []
for pt in allVert:
        point.X = pt[2]
        point.Y = pt[3]
        pointGeometry = arcpy.PointGeometry(point)
        ptGeom.append(pointGeometry)

## create geometry in feature class
arcpy.CopyFeatures_management(ptGeom, pt3D)

del ptGeom
del point

## add fields to populate with line ID, x and y coords 
arcpy.AddField_management(pt3D, "ORIG_FID", "LONG", 5)
arcpy.AddField_management(pt3D, "POINT_X", "DOUBLE", 12, 4)
arcpy.AddField_management(pt3D, "POINT_Y", "DOUBLE", 12, 4)

## append attributes
cur = arcpy.UpdateCursor(pt3D)
count = 1
for row in cur:
        if row.getValue('OBJECTID') == count:
                row.setValue('ORIG_FID', allVert[count-1][1])
                row.setValue('POINT_X', allVert[count-1][2])
                row.setValue('POINT_Y', allVert[count-1][3])        
                cur.updateRow(row)
                count = count + 1

## get DEM values from point intersection
DEMin = Raster("C:\\ArcGIS\\MOI\\MOI_scratch.gdb\\INPUT_DEM")

## create output points ( with x, y and z)
try:
        arcpy.DeleteFeatures_management("C:\\ArcGIS\\MOI\\MOI_template.gdb\\points3D")
except:
        print ""
        
pt3Dout = arcpy.CreateFeatureclass_management("C:\\ArcGIS\\MOI\\MOI_template.gdb\\", "points3D", "POINT", has_m="DISABLED", has_z="ENABLED", spatial_reference=lineFC)

## Execute ExtractValuesToPoints
ExtractValuesToPoints(pt3D, DEMin, pt3Dout, "NONE", "VALUE_ONLY")

addLayer = arcpy.mapping.Layer("C:\\ArcGIS\\MOI\\MOI_template.gdb\\points3D")
arcpy.mapping.AddLayer(df, addLayer)
arcpy.MakeFeatureLayer_management(pt3Dout, "point3D_lyr")
arcpy.ApplySymbologyFromLayer_management(addLayer, "C:/ArcGIS/MOI/3D_points.lyr")
arcpy.RefreshActiveView()
arcpy.RefreshTOC() 

## remove temp data
arcpy.Delete_management("C:\\ArcGIS\\MOI\\MOI_scratch.gdb\\temp3D")

arcpy.AddMessage("3D points loaded")

## ----------------------------------

arcpy.AddMessage("Processing MOI ... ")

## run MOI estimation
def PointsToPlane():

        Xm = []
        Ym = []
        Zm = []

        Xdiff = []
        Ydiff = []
        Zdiff = []

        out = []

## create table of values 
        MainDict = {}
        InFC = "C:\\ArcGIS\\MOI\\MOI_template.gdb\\points3D"
        outFC = "C:\\ArcGIS\\MOI\\MOI_template.gdb\\GeoLines"
        
        searchRows = arcpy.SearchCursor(InFC)

        for rec in searchRows:
                idValue = rec.getValue("ORIG_FID")
                xValue = rec.getValue("POINT_X")
                yValue = rec.getValue("POINT_Y")
                zValue = rec.getValue("RASTERVALU")

                if not idValue in MainDict:
                        MainDict[idValue] = Record(idValue)

                currentRec = MainDict[idValue]
                currentRec.coords.append(Coord(xValue, yValue, zValue))
                
        print "LINE ID = ", MainDict.keys()

        
        for recKey in MainDict.keys():
                rec = MainDict[recKey]
                rec.average()
                rec.plane_calc()
                rec.RMSE()
                rec.orientation()
                ## write to feature class
                outDat = rec.writer(outFC)
                out.append((outDat,rec.Error()))

        return(out)

## get outputs for additional analysis
Output = PointsToPlane()

Dat = []
for i in Output:
        Dat.append(i[0])

ptErr = []
for i in Output:
        for j in i[1]:
                ptErr.append(j)

## add pt error (to plane)
arcpy.AddField_management(pt3Dout, "Error", "DOUBLE", 12, 4)

count = 1
cur = arcpy.UpdateCursor(pt3Dout)
for row in cur:
        if row.getValue("OBJECTID")==count:
                row.setValue("Error",ptErr[count-1])
                cur.updateRow(row)
                count = count + 1

arcpy.AddMessage("MOI completed")
             
## ----------------------------------

arcpy.AddMessage("Creating mid points ... ")

## create midpoint feature class
try:
        arcpy.Delete_management("C:\\ArcGIS\\MOI\\MOI_template.gdb\\GeoPoint")
except:
        print ""

midpt = arcpy.CreateFeatureclass_management("C:\\ArcGIS\\MOI\\MOI_template.gdb\\", "GeoPoint", "POINT", has_m="DISABLED", has_z="ENABLED", spatial_reference=lineFC)

## create new points
point = arcpy.Point()
ptGeom = []
for pt in midPoints:
        point.X = pt[1]
        point.Y = pt[2]
        pointGeometry = arcpy.PointGeometry(point)
        ptGeom.append(pointGeometry)

## create geometry in feature class
arcpy.CopyFeatures_management(ptGeom, midpt)

del ptGeom
del point

## get list of fields to add to midpt
fieldList = [f.name for f in arcpy.ListFields(lineFC)]

## add fields to populate with line ID
for name in fieldList[2:16]:
        if name == "ID":
                arcpy.AddField_management(midpt, name, "LONG", 5)
        elif name == "N_points":
                arcpy.AddField_management(midpt, name, "LONG", 5)
        elif name == "Dip":
                arcpy.AddField_management(midpt, name, "LONG", 5)
        elif name == "Dipdir":
                arcpy.AddField_management(midpt, name, "LONG", 5)
        else:
                arcpy.AddField_management(midpt, name, "DOUBLE", 20, 4)
                
arcpy.AddField_management(midpt, "eigen1", "DOUBLE", 20, 4)
arcpy.AddField_management(midpt, "eigen2", "DOUBLE", 20, 4)
arcpy.AddField_management(midpt, "eigen3", "DOUBLE", 20, 4)
arcpy.AddField_management(midpt, "Length", "DOUBLE", 20, 4)

cur =  arcpy.SearchCursor(lineFC)
leng = []
for row in cur:
        leng.append(row.getValue('SHAPE_Length'))

## append attributes
cur = arcpy.UpdateCursor(midpt)
count = 1
for row in cur:
        if row.getValue('OBJECTID') == count:
                row.setValue('ID', Dat[count-1][0])
                row.setValue('N_points',Dat[count-1][1])
                row.setValue('Xmean', Dat[count-1][2])
                row.setValue('Ymean', Dat[count-1][3])
                row.setValue('Zmean', Dat[count-1][4])
                row.setValue('M', Dat[count-1][5])
                row.setValue('K', Dat[count-1][6])
                row.setValue('RMSE', Dat[count-1][7])
                row.setValue('A', Dat[count-1][8])
                row.setValue('B', Dat[count-1][9])
                row.setValue('C', Dat[count-1][10])
                row.setValue('D', Dat[count-1][11])
                row.setValue('Dip', Dat[count-1][12])
                row.setValue('Dipdir', Dat[count-1][13])
                row.setValue('eigen1', Dat[count-1][14])
                row.setValue('eigen2', Dat[count-1][15])
                row.setValue('eigen3', Dat[count-1][16])  
                row.setValue('Length', leng[count-1])
                cur.updateRow(row)
                count = count + 1

## add points as layer
addLayer = arcpy.mapping.Layer("C:/ArcGIS/MOI/MOI_template.gdb/GeoPoint")
arcpy.mapping.AddLayer(df, addLayer)
# get symbology
arcpy.MakeFeatureLayer_management("C:/ArcGIS/MOI/MOI_template.gdb/GeoPoint", "GeoPoint_lyr")
arcpy.ApplySymbologyFromLayer_management(addLayer, "C:/ArcGIS/MOI/GeoPoint.lyr")
# label points
addLayer.labelClasses[0].expression = "[Dip]"
addLayer.showLabels = True
# refresh view
arcpy.RefreshActiveView()
arcpy.RefreshTOC() 

arcpy.AddMessage("Mid points created")

## ----------------------------------

## write to csv
arcpy.AddMessage("Writing data to MOI_dat.csv")

fieldList = [f.name for f in arcpy.ListFields(midpt)][2:20]

with open('C:/ArcGIS/MOI/MOI_dat.csv', 'wb') as csvfile:
        datwriter = csv.writer(csvfile, delimiter=',')
        datwriter.writerow(fieldList)
        count = 1
        for row in Dat:
                a,b,c,d,e,f,g,h,i,j,k,l,m,n,o,p,q = row
                datwriter.writerow([a,b,c,d,e,f,g,h,i,j,k,l,m,n,o,p,q,leng[count-1]])
                count = count + 1

## ----------------------------------

## plot evaluation measures
arcpy.AddMessage("Plotting M-K")

## basic plot parameters
moiPNG = "C:/ArcGIS/MOI/MOI.png"

fig = plt.figure()
plt.axis([0, 7, 0, 7])

## set up lines
mLinesX = [2,0]
mLinesY = [0,2]

for i in [1,2,3]:
        plt.plot([x * i for x in mLinesX] ,[x * i for x in mLinesY], color="gray")
           
for k in [0.2,0.5,1,2,5]:
        ky = k * 7
        plt.plot([0,7],[0,ky], color="gray")

## set up "good fit" polygon
ax = fig.add_subplot(111)
Path = mpath.Path
code = [Path.MOVETO,Path.LINETO,Path.LINETO,Path.LINETO,Path.CLOSEPOLY]
verts = [(4,0),(2.222,2.222*0.8),(7,0.8*7),(7,0),(4,0)] 
path = mpath.Path(verts, code)
patch = mpatch.PathPatch(path, facecolor='green', alpha=0.5)
ax.add_patch(patch)

## get point coords and plot
ptCoord = []
for i in Dat:
        x = np.log(i[15]/i[16])
        y = np.log(i[14]/i[15])
        plt.plot(x,y, 'bo')
        ax.text(x,y, str(i[0]),horizontalalignment='left',verticalalignment='bottom')

## axis labels
plt.xlabel('ln(EigenVal2/EigenVal3)')
plt.ylabel('ln(EigenVal1/EigenVal2)')
plt.title('MOI fit')

## annotate plot
ax.text(1, 1.1, 'M = 2',horizontalalignment='center',rotation=-45,size="smaller")
ax.text(2, 2.1, 'M = 4',horizontalalignment='center',rotation=-45,size="smaller")
ax.text(3, 3.1, 'M = 6',horizontalalignment='center',rotation=-45,size="smaller")
ax.text(1.2, 6, 'K = 5',horizontalalignment='center',rotation=73,size="smaller")
ax.text(3, 6.1, 'K = 2',horizontalalignment='center',rotation=56,size="smaller")
ax.text(5, 5.1, 'K = 1',horizontalalignment='center',rotation=42,size="smaller")
ax.text(6, 3.1, 'K = 0.5',horizontalalignment='center',rotation=24,size="smaller")
ax.text(6, 1.2, 'K = 0.2',horizontalalignment='center',rotation=10,size="smaller")

## save and open plot for display
fig.savefig(moiPNG, dpi=200)
os.startfile(moiPNG)

## end
## ---------------------------------------------------------------------------
