import math
import numpy as np
from obspy.taup import TauPyModel
from obspy.taup.velocity_model import VelocityModel
from obspy.clients.iris import Client

def sind(deg):
    rad = math.radians(deg)
    return math.sin(rad)


def cosd(deg):
    rad = math.radians(deg)
    return math.cos(rad)


def tand(deg):
    rad = math.radians(deg)
    return math.tan(rad)


def cotd(deg):
    rad = math.radians(deg)
    return math.cos(rad) / math.sin(rad)


def asind(x):
    rad = math.asin(x)
    return math.degrees(rad)


def acosd(x):
    rad = math.acos(x)
    return math.degrees(rad)


def atand(x):
    rad = math.atan(x)
    return math.degrees(rad)


def km2deg(kilometers):
    return kilometers / 111.19


def deg2km(degree):
    return degree * 111.19


class distaz:
    """
    c Subroutine to calculate the Great Circle Arc distance
    c    between two sets of geographic coordinates
    c
    c Equations take from Bullen, pages 154, 155
    c
    c T. Owens, September 19, 1991
    c           Sept. 25 -- fixed az and baz calculations
    c
    P. Crotwell, Setember 27, 1995
    Converted to c to fix annoying problem of fortran giving wrong
    answers if the input doesn't contain a decimal point.
    
    H. P. Crotwell, September 18, 1997
    Java version for direct use in java programs.
    *
    * C. Groves, May 4, 2004
    * Added enough convenience constructors to choke a horse and made public double
    * values use accessors so we can use this class as an immutable

    H.P. Crotwell, May 31, 2006
    Port to python, thus adding to the great list of languages to which
    distaz has been ported from the origin fortran: C, Tcl, Java and now python
    and I vaguely remember a perl port. Long live distaz! 
    """

    def __init__(self, lat1, lon1, lat2, lon2):
        
        self.stalat = lat1
        self.stalon = lon1
        self.evtlat = lat2
        self.evtlon = lon2
        '''
        if (lat1 == lat2) and (lon1 == lon2):
            self.delta = 0.0
            self.az = 0.0
            self.baz = 0.0
            return
        '''

        rad = 2. * math.pi / 360.0
        """
        c
        c scolat and ecolat are the geocentric colatitudes
        c as defined by Richter (pg. 318)
        c
        c Earth Flattening of 1/298.257 take from Bott (pg. 3)
        c
        """
        sph = 1.0 / 298.257

        scolat = math.pi / 2.0 - np.arctan((1. - sph) * (1. - sph) * np.tan(lat1 * rad))
        ecolat = math.pi / 2.0 - np.arctan((1. - sph) * (1. - sph) * np.tan(lat2 * rad))
        slon = lon1 * rad
        elon = lon2 * rad
        """
	c
	c  a - e are as defined by Bullen (pg. 154, Sec 10.2)
	c     These are defined for the pt. 1
	c
        """
        a = np.sin(scolat) * np.cos(slon)
        b = np.sin(scolat) * np.sin(slon)
        c = np.cos(scolat)
        d = np.sin(slon)
        e = -np.cos(slon)
        g = -c * e
        h = c * d
        k = -np.sin(scolat)
        """
	c
	c  aa - ee are the same as a - e, except for pt. 2
	c
        """
        aa = np.sin(ecolat) * np.cos(elon)
        bb = np.sin(ecolat) * np.sin(elon)
        cc = np.cos(ecolat)
        dd = np.sin(elon)
        ee = -np.cos(elon)
        gg = -cc * ee
        hh = cc * dd
        kk = -np.sin(ecolat)
        """
	c
	c  Bullen, Sec 10.2, eqn. 4
	c
        """
        delrad = np.arccos(a * aa + b * bb + c * cc)
        self.delta = delrad / rad
        """
	c
	c  Bullen, Sec 10.2, eqn 7 / eqn 8
	c
	c    pt. 1 is unprimed, so this is technically the baz
	c
	c  Calculate baz this way to avoid quadrant problems
	c
        """
        rhs1 = (aa - d) * (aa - d) + (bb - e) * (bb - e) + cc * cc - 2.
        rhs2 = (aa - g) * (aa - g) + (bb - h) * (bb - h) + (cc - k) * (cc - k) - 2.
        dbaz = np.arctan2(rhs1, rhs2)

        dbaz_idx = np.where(dbaz < 0.0)[0]
        if len(dbaz_idx) != 0:
            if isinstance(dbaz, (int, float)):
                dbaz += 2 * math.pi
            else:
                dbaz[dbaz_idx] += 2 * math.pi

        self.baz = dbaz / rad
        """
	c
	c  Bullen, Sec 10.2, eqn 7 / eqn 8
	c
	c    pt. 2 is unprimed, so this is technically the az
	c
	"""
        rhs1 = (a - dd) * (a - dd) + (b - ee) * (b - ee) + c * c - 2.
        rhs2 = (a - gg) * (a - gg) + (b - hh) * (b - hh) + (c - kk) * (c - kk) - 2.
        daz = np.arctan2(rhs1, rhs2)

        daz_idx = np.where(daz < 0.0)[0]
        if len(daz_idx) != 0:
            if isinstance(daz, (int, float)):
                daz += 2 * math.pi
            else:
                daz[daz_idx] += 2 * math.pi

        self.az = daz / rad
        """
	c
	c   Make sure 0.0 is always 0.0, not 360.
	c
	"""
        idx = np.where(np.abs(self.baz - 360.) < .00001)[0]
        if len(idx) != 0:
            if isinstance(self.baz, float):
                self.baz = 0.0
            else:
                self.baz[idx] = 0.0
        idx = np.where(np.abs(self.baz) < .00001)[0]
        if len(idx) != 0:
            if isinstance(self.baz, float):
                self.baz = 0.0
            else:
                self.baz[idx] = 0.0

        idx = np.where(np.abs(self.az - 360.) < .00001)[0]
        if len(idx) != 0:
            if isinstance(self.az, float):
                self.az = 0.0
            else:
                self.az[idx] = 0.0
        idx = np.where(np.abs(self.az) < .00001)[0]
        if len(idx) != 0:
            if isinstance(self.az, float):
                self.az = 0.0
            else:
                self.az[idx] = 0.0
        
        la_idx = np.where(lat1 == lat2)[0]
        lo_idx = np.where(lon1 == lon2)[0]
        idx = la_idx[np.where(la_idx == lo_idx)[0]]
        if len(idx) != 0:
            if isinstance(self.delta, float):
                self.delta = 0.
            else:
                self.delta[idx] = 0.
            if isinstance(self.az, float):
                self.az = 0.
            else:
                self.az[idx] = 0.
            if isinstance(self.baz, float):
                self.baz = 0.
            else:
                self.baz[idx] = 0.

    def getDelta(self):
        return self.delta

    def getAz(self):
        return self.az

    def getBaz(self):
        return self.baz

    def degreesToKilometers(self):
        return self.delta * 111.19
 
def loaddt(fna):
    ''' Loading the time residual'''
    global evt, nsta, sta
    evt = []
    f = open(fna)
    l1 = f.readline()
    for item in l1.split():
        evt = np.append(evt,float(item))
        
    nsta = int(f.readline())
    sta = np.zeros((nsta,3))
    for i in range(nsta):
        line = f.readline()
        sta[i,0],sta[i,1],sta[i,2] = float(line.split()[0]),float(line.split()[1]),float(line.split()[2])
    
def dt(x):
    # Store the list of function calls
    # obspy theoretical traveltime 
    #
    ###client = Client()
    TT = np.zeros(nsta)
    for i in range(nsta):
        ###delta = client.distaz(stalat=sta[i,0], stalon=sta[i,1],\
        ###    evtlat=evt[0], evtlon=evt[1])['distance']
        dist = distaz(x[0], x[1], sta[i,0], sta[i,1])
        delta = dist.delta
        arrivals  = m.get_travel_times(source_depth_in_km=x[2],distance_in_degree=delta,\
            phase_list=ph, receiver_depth_in_km=0.0)
        TT[i] = arrivals[1].time- arrivals[0].time 
    return np.sqrt( sum( ( TT[:]-sta[:,2] )**2) )
# 
# reading data 
#
inf =  "Bolivia_dt.dat" 
loaddt(inf)
#  
global m, ph
ph = ["PKIKP","pPKIKP"]
modna='ak135'
m = TauPyModel(model=modna)
dd = 30.0
dla = 0.3
dlo = 0.3
dep0 = evt[2] - dd
dep1 = evt[2] + dd
la0 = evt[0] - dla
la1 = evt[0] + dla
lo0 = evt[1] - dlo
lo1 = evt[1] + dlo

dep = np.random.uniform(dep0, dep1)
la_v = np.random.uniform(la0, la1)
lo_v = np.random.uniform(lo0, lo1)
x = [la_v,lo_v,dep]
#x = evt
print(x,dt(x))
T = 100. # inital T
std = 0.00000001
a = 0.99    # attenuation
i=0
x_new = x
y = dt(x)
while T > std:
    i = i + 1
    x_new[0] = x[0] + 0.02*np.random.uniform(-1.0, 1.0)
    x_new[1] = x[1] + 0.02*np.random.uniform(-1.0, 1.0)
    #x_new[0] = x[0] + 0.01*np.random.uniform(-1.0, 1.0)
    #x_new[1] = x[1] + 0.01*np.random.uniform(-1.0, 1.0)
    x_new[2] = x[2] + np.random.uniform(-1.0, 1.0)
    
    if dep0 <= x_new[2] <= dep1 and la0 <= x_new[0] <= la1 and lo0 <= x_new[1] <= lo1 :
        y_new = dt(x_new)
        #print(i,x_new,y_new)
        if y_new < y:
            x = x_new
            y = y_new
            print(i,x_new,y_new)
        else:
            p = np.exp( (y-y_new)/T )
            r = np.random.uniform(0,1)
            if p > r:
                x = x_new
    T = T * a
print("======================================")    
print(i,x,dt(x))  
print( evt, dt(evt) )  