import warnings
warnings.filterwarnings("ignore")

# core python 
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm
import ipywidgets
import os

# tools in the simPEG Ecosystem 
import discretize  # for creating computational meshes

# linear solvers
try: 
    from pymatsolver import Pardiso as Solver  # this is a fast linear solver 
except ImportError:
    from SimPEG import SolverLU as Solver  # this will be slower

# SimPEG inversion machinery
from SimPEG import (
    Data, maps,
    data_misfit, regularization, optimization, inverse_problem, 
    inversion, directives
) 

# DC resistivity and IP modules
from SimPEG.electromagnetics import resistivity as dc
from SimPEG.electromagnetics import induced_polarization as ip

# set the font size in the plots
from matplotlib import rcParams
rcParams["font.size"] = 14

os.listdir('century')

line = "46800E"
print(os.listdir(os.path.join('century',line)))

dc_data_file = f"./century/{line}/{line[:-1]}POT.OBS"

Comment line: description of data
   n_sources dipole(1) dipole(1) 
     A_location  B_location   n_receivers
   M_location  N_location  measured_potential  standard_deviation
   M_location  N_location  measured_potential  standard_deviation
     A_location  B_location   n_receivers
   M_location  N_location  measured_potential  standard_deviation
   M_location  N_location  measured_potential  standard_deviation
   ....

with open(dc_data_file) as f:
    print(f.read())

def read_dcip_data(filename, verbose=True):
    """
    Read in a .OBS file from the Century data set into a python dictionary. 
    The format is the old UBC-GIF DCIP format.
    
    Parameters
    ----------
    filename : str
        Path to the file to be parsed
    
    verbose: bool
        Print some things? 
    
    
    Returns
    -------
    dict
        A dictionary with the locations of
        - a_locations: the positive source electrode locations (numpy array) 
        - b_locations: the negative source electrode locations (numpy array) 
        - m_locations: the receiver locations (list of numpy arrays)
        - n_locations: the receiver locations (list of numpy arrays)
        - observed_data: observed data (list of numpy arrays)
        - standard_deviations: assigned standard deviations (list of numpy arrays)
        - n_sources: number of sources (int)
    
    """
    
    # read in the text file as a numpy array of strings (each row is an entry)
    contents = np.genfromtxt(filename, delimiter=' \n', dtype=np.str)
    
    # the second line has the number of sources, current, and data type (voltages if 1)
    n_sources = int(contents[1].split()[0])
    
    if verbose is True: 
        print(f"number of sources: {n_sources}")
    
    # initialize storage for the electrode locations and data
    a_locations = np.zeros(n_sources)
    b_locations = np.zeros(n_sources)
    m_locations = []
    n_locations = []
    observed_data = []
    standard_deviations = []
    
    # index to track where we have read in content 
    content_index = 1 
    
    # loop over sources 
    for i in range(n_sources):
        # start by reading in the source info 
        content_index = content_index + 1  # read the next line
        a_location, b_location, nrx = contents[content_index].split()  # this is a string

        # convert the strings to a float for locations and an int for the number of receivers
        a_locations[i] = float(a_location)
        b_locations[i] = float(b_location)
        nrx = int(nrx)

        if verbose is True: 
            print(f"Source {i}: A-loc: {a_location}, B-loc: {b_location}, N receivers: {nrx}")

        # initialize space for receiver locations, observed data associated with this source
        m_locations_i, n_locations_i = np.zeros(nrx), np.zeros(nrx)
        observed_data_i, standard_deviations_i = np.zeros(nrx), np.zeros(nrx)

        # read in the receiver info 
        for j in range(nrx):
            content_index = content_index + 1  # read the next line
            m_location, n_location, datum, std = contents[content_index].split()

            # convert the locations and data to floats, and store them
            m_locations_i[j] = float(m_location)
            n_locations_i[j] = float(n_location)
            observed_data_i[j] = float(datum)
            standard_deviations_i[j] = float(std)

        # append the receiver info to the lists
        m_locations.append(m_locations_i)
        n_locations.append(n_locations_i)
        observed_data.append(observed_data_i)
        standard_deviations.append(standard_deviations_i)
    
    return {
        "a_locations": a_locations,
        "b_locations": b_locations, 
        "m_locations": m_locations,
        "n_locations": n_locations,
        "observed_data": observed_data, 
        "standard_deviations": standard_deviations,
        "n_sources": n_sources, 
    }

dc_data_dict = read_dcip_data(dc_data_file)

for key, value in dc_data_dict.items():
    print(f"{key:<20}: {type(value)}")

dc_data_dict["observed_data"]

m_locs = np.array([
    [26700.0, 0.],
    [26800.0, 0.]
])

n_locs = np.array([
    [26800.0, 0.],
    [26900.0, 0.]
])

rx = dc.receivers.Dipole(locations_m=m_locs, locations_n=n_locs) 
rx_list = [rx]

a_loc = np.r_[26000.0, 0.]
b_loc = np.r_[26100.0, 0.]

src = dc.sources.Dipole(location_a=a_loc, location_b=b_loc, receiver_list=rx_list)

survey = dc.Survey(source_list=[src1, src2, ...])

# initialize an empty list for each 
source_list = []

for i in range(dc_data_dict["n_sources"]):
    
    # receiver electrode locations in 2D 
    m_locs = np.vstack([
        dc_data_dict["m_locations"][i], 
        np.zeros_like(dc_data_dict["m_locations"][i])
    ]).T
    n_locs = np.vstack([
        dc_data_dict["n_locations"][i],
        np.zeros_like(dc_data_dict["n_locations"][i])
    ]).T
    
    # construct the receiver object 
    receivers = dc.receivers.Dipole(locations_m=m_locs, locations_n=n_locs)
    
    # construct the source 
    source = dc.sources.Dipole(
        location_a=np.r_[dc_data_dict["a_locations"][i], 0.],
        location_b=np.r_[dc_data_dict["b_locations"][i], 0.],
        receiver_list=[receivers]
    )
    
    # append the new source to the source list
    source_list.append(source)

src1 = source_list[0]

src1.receiver_list

survey = dc.Survey(source_list=source_list)

survey.nD

dobs = [d(source1, receiver1), d(source1, receiver2), d(source2, receiver1), ...]

dc_data = Data(
    survey=survey, 
    dobs=np.hstack(dc_data_dict["observed_data"]),
    standard_deviation=np.hstack(dc_data_dict["standard_deviations"])
)

# plot psuedosection
fig, ax = plt.subplots(1, 1, figsize=(12, 4))
dc.utils.plot_pseudosection(
    dc_data, data_type="potential",
    plot_type="contourf", data_location=True, ax=ax,
)
ax.set_aspect(1.5)  # some vertical exxageration

def plot_building_pseudosection(source_ind=0):
    """
    A plotting function to visualize how a pseudosection is built up. 
    """
    
    fig, ax = plt.subplots(1, 1, figsize=(12, 4))
    mid_x, mid_z = dc.utils.source_receiver_midpoints(dc_data.survey)
    ax.plot(mid_x, mid_z, '.k', alpha=0.5)
    
    # plot the source location
    a_loc = dc_data_dict["a_locations"][source_ind]
    b_loc = dc_data_dict["b_locations"][source_ind]
    src_mid = (a_loc+b_loc)/2
    ax.plot(np.r_[a_loc, b_loc], np.r_[0, 0], 'C0.')
    ax.plot(src_mid, np.r_[0], 'C0x', ms=6)

    # plot the receiver locations
    m_locs = dc_data_dict["m_locations"][source_ind]
    n_locs = dc_data_dict["n_locations"][source_ind]
    rx_mid = (m_locs+n_locs)/2
    ax.plot(np.r_[m_locs, n_locs], np.zeros(2*len(m_locs)), 'C1.')
    ax.plot(rx_mid, np.zeros_like(m_locs), 'C1x', ms=6)
    
    # plot where the pseudosection points should be 
    pseudo_x = (rx_mid + src_mid)/2.
    pseudo_z = -np.abs(rx_mid-src_mid)/2.
    
    ax.plot(np.r_[src_mid, pseudo_x], np.r_[0, pseudo_z], '-k', alpha=0.3)
    
    # draw lines to the points we are plotting
    for rx_x, px, pz in zip(rx_mid, pseudo_x, pseudo_z): 
        ax.plot(np.r_[px, rx_x], np.r_[pz, 0], '-k', alpha=0.3)
    
    # these are the data corresponding to the source-receiver pairs we are looking at
    ax.plot(pseudo_x, pseudo_z, 'C2o')
    ax.set_xlabel("Northing (m)")
    ax.set_ylabel("n-spacing")
    
    ax.set_yticklabels([])
    ax.set_xlim([25900, 29300])
    ax.set_aspect(1.5) 
    

ipywidgets.interact(
    plot_building_pseudosection, 
    source_ind=ipywidgets.IntSlider(min=0, max=int(dc_data_dict["n_sources"])-1, value=0)
)

# plot psuedosection
fig, ax = plt.subplots(1, 1, figsize=(12, 4))
dc.utils.plot_pseudosection(
    dc_data, data_type="apparent resistivity", 
    plot_type="contourf", data_location=True, ax=ax, cbar_opts={"pad":0.25}
)
ax.set_aspect(1.5)  # some vertical exxageration
ax.set_title(f"DC: {line} Pseudosection")
ax.set_xlabel("Northing (m)")

apparent_resistivity = dc.utils.apparent_resistivity(dc_data)

fig, ax = plt.subplots(1, 1)
out = ax.hist(np.log10(apparent_resistivity), bins=30)
ax.set_xlabel("log$_{10}(\\rho_a)$")

# rho0 = 10**np.mean(np.log10(apparent_resistivity))
rho0 = np.median(apparent_resistivity)
rho0

min_electrode_spacing = np.min(np.abs(survey.locations_a[:, 0] - survey.locations_b[:, 0]))
n_cells_per_spacing = 4

dx = min_electrode_spacing / n_cells_per_spacing
dz = min_electrode_spacing / n_cells_per_spacing

print(
    f"We will use a base mesh with {dx} m x {dz} m cells"
)

core_domain_x = np.r_[
    survey.electrode_locations[:, 0].min(),
    survey.electrode_locations[:, 0].max()
]

core_domain_x

# find the maximum spacing between source, receiver midpoints
mid_ab = (survey.locations_a + survey.locations_b)/2
mid_mn = (survey.locations_m + survey.locations_n)/2
separation_ab_mn = np.abs(mid_ab - mid_mn)
max_separation = separation_ab_mn.max()
print(f"The maximum separation is {max_separation}m")

core_domain_z = np.r_[-max_separation/3., 0.]
core_domain_z

n_core_extra = 4 
n_core_x = np.ceil(np.diff(core_domain_x)/dx) + n_core_extra*2  # on each side
n_core_z = np.ceil(np.diff(core_domain_z)/dz) + n_core_extra  # just below

padding_factor = 1.3
n_pad_x = 10
n_pad_z = 10

hx = [(dx, n_pad_x, -padding_factor), (dx, n_core_x), (dx, n_pad_x, padding_factor)]
hz = [(dz, n_pad_z, -padding_factor), (dz, n_core_z)]

mesh = discretize.TensorMesh([hx, hz])

# origin of the mesh
mesh.x0 = np.r_[
    -mesh.hx[:n_pad_x + n_core_extra].sum() + core_domain_x.min(),
    -mesh.hy.sum()
]

mesh

fig, ax = plt.subplots(1, 1, figsize=(12, 4))
mesh.plotGrid(ax=ax)
ax.set_aspect(1.5)

# forward simulation 
# Use Exponential Map: m = log(rho)
mapping = maps.ExpMap(mesh)

# Generate 2.5D DC problem
simulation_dc = dc.Simulation2DNodal(
    mesh, rhoMap=mapping, solver=Solver, survey=survey, storeJ=True
)

%%time 

# run the forward simulation over the half-space & plot apparent resistivities
m0 = np.ones(mesh.nC) * np.log(rho0)
d0 = simulation_dc.make_synthetic_data(m0)

# plot psuedosection
fig, ax = plt.subplots(1, 1, figsize=(12, 4))

clim = np.r_[apparent_resistivity.min(), apparent_resistivity.max()]
dc.utils.plot_pseudosection(
    d0, data_type="apparent resistivity",  clim=clim,
    plot_type="contourf", data_location=True, ax=ax, 
    cbar_opts={"pad":0.25}
)
ax.set_aspect(1.5)  # some vertical exxageration
ax.set_title(f"{line} Pseudosection")
ax.set_xlabel("Northing (m)")

apparent_resistivity_m0 = dc.utils.apparent_resistivity(d0)
percent_error = (apparent_resistivity_m0 - rho0)/rho0*100

fig, ax = plt.subplots(1, 1)
ax.plot(separation_ab_mn[:, 0], percent_error, '.', alpha=0.4)
ax.set_xlabel("AB - MN separation")
ax.set_ylabel("percent difference (%)")
ax.set_title("Error in apparent resistivities")
ax.set_ylim(7*np.r_[-1, 1])
ax.grid(alpha=0.3)

# set up the data misfit
dmisfit = data_misfit.L2DataMisfit(data=dc_data, simulation=simulation_dc)

# Related to inversion
reg = regularization.Tikhonov(
    mesh, 
    alpha_s=1./mesh.hx.min()**2,  # we will talk about these choices in notebook 3 on inversion
    alpha_x=1, 
    alpha_y=1,  # since this is a 2D problem, the second dimension is "y" 
)

opt = optimization.InexactGaussNewton(maxIter=20, maxIterCG=20)
inv_prob = inverse_problem.BaseInvProblem(dmisfit, reg, opt)

# directives
beta_est = directives.BetaEstimate_ByEig(beta0_ratio=1e0)
beta_schedule = directives.BetaSchedule(coolingFactor=4, coolingRate=2)
target = directives.TargetMisfit()

inv = inversion.BaseInversion(
    inv_prob, directiveList=[beta_est, beta_schedule, target]
)

# Run inversion
mopt = inv.run(m0)

# plot psuedosection
fig, ax = plt.subplots(2, 1, figsize=(12, 8))

clim = np.r_[apparent_resistivity.min(), apparent_resistivity.max()]

for a, data, title in zip(
    ax,
    [dc_data.dobs, inv_prob.dpred],
    ["true", "recovered"]
): 
    dc.utils.plot_pseudosection(
        dc_data, dobs=data, data_type="apparent resistivity",
        plot_type="contourf", data_location=True, ax=a, 
        cbar_opts={"pad":0.25}, clim=clim
    )
    a.set_title(title)


for a in ax: 
    a.set_aspect(1.5)  # some vertical exxageration
    a.set_xlabel("Northing (m)")
    a.set_ylabel("Pseudo-depth")

plt.tight_layout()

# plot misfit
fig, ax = plt.subplots(1, 1, figsize=(12, 4))

normalized_misfit = (inv_prob.dpred - dc_data.dobs) / dc_data.standard_deviation

dc.utils.plot_pseudosection(
    dc_data, dobs=normalized_misfit, data_type="misfit",
    plot_type="contourf", data_location=True, ax=ax, 
    cbar_opts={"pad":0.25}
)

ax.set_title("DC misfit")
ax.set_aspect(1.5)  # some vertical exxageration
ax.set_xlabel("Northing (m)")
ax.set_ylabel("n-spacing")

cb_axes = plt.gcf().get_axes()[-1]
cb_axes.set_xlabel('normalized misfit')

plt.tight_layout()

# run inversion, plot tikhonov curves

def load_leapfrog_geologic_section(filename="./century/geologic_section.csv"):
    """
    Load the geologic cross section. 
    """
    fid = open(filename, 'r')
    lines = fid.readlines()
    data = []
    data_tmp = []
    for line in lines[2:]:
        line_data = (line.split(',')[:3])
        if 'End' in line:
            data.append(np.vstack(data_tmp)[:,[0, 2]])
            data_tmp = []
        else:
            data_tmp.append(np.array(line_data, dtype=float))
    return data

geologic_section = load_leapfrog_geologic_section()

fig, ax = plt.subplots(1,1, figsize=(10, 5))
rho = (mapping*mopt)
out = mesh.plotImage(
    rho, pcolorOpts={'norm':LogNorm(), 'cmap':'Spectral'}, ax=ax,
    clim=(30, 1000)
)
for data in geologic_section:
    ax.plot(data[:,0], data[:,1], 'k--', alpha=0.5)
ax.set_xlim(core_domain_x)
ax.set_ylim((-400, 0))
cb = plt.colorbar(out[0], fraction=0.05, orientation='horizontal', ax=ax, pad=0.2)
cb.set_label("Resistivity ($\Omega$m)")
ax.set_xlabel('Easting (m)')
ax.set_ylabel('Elevation (m)')
ax.set_aspect(1.5)  # some vertical exxageration
plt.tight_layout()

mesh_ubc = discretize.TensorMesh.readUBC("century/46800E/468MESH.DAT")

def read_ubc_model(filename, mesh_ubc=mesh_ubc): 
    """
    A function to read a UBC conductivity or chargeability model. 
    """
    values = np.genfromtxt(
        filename, delimiter=' \n',
        dtype=str, comments='!', skip_header=1
    )
    tmp = np.hstack([np.array(value.split(), dtype=float) for value in values])
    model_ubc = discretize.utils.mkvc(tmp.reshape(mesh_ubc.vnC, order='F')[:,::-1])
    return model_ubc
    

sigma_ubc = read_ubc_model("century/46800E/DCMODA.CON")
rho_ubc = 1./sigma_ubc

fig, ax = plt.subplots(2, 1, figsize=(10, 8))
rho = (mapping*mopt)
clim = (30, 1000)

out = mesh.plotImage(
    rho, pcolorOpts={'norm':LogNorm(), 'cmap':'Spectral'}, ax=ax[0],
    clim=clim 
)
cb = plt.colorbar(out[0], fraction=0.05, orientation='horizontal', ax=ax[0], pad=0.2)
cb.set_label("Resistivity ($\Omega$m)")
ax[0].set_title("SimPEG")

out = mesh_ubc.plotImage(
    rho_ubc, pcolorOpts={'norm':LogNorm(), 'cmap':'Spectral'}, ax=ax[1],
    clim=clim 
)
cb = plt.colorbar(out[0], fraction=0.05, orientation='horizontal', ax=ax[1], pad=0.2)
cb.set_label("Resistivity ($\Omega$m)")
ax[1].set_title("UBC")

for a in ax: 
    for data in geologic_section:
        a.plot(data[:,0], data[:,1], 'k--', alpha=0.5)
    a.set_xlim(core_domain_x)
    a.set_ylim((-400, 0))
    a.set_aspect(1.5)  # some vertical exxageration
    a.set_xlabel('Easting (m)')
    a.set_ylabel('Elevation (m)')

plt.tight_layout()

ip_data_file = f"./century/{line}/{line[:-1]}IP.OBS"

ip_data_dict = read_dcip_data(ip_data_file, verbose=False)

# initialize an empty list for each 
source_list_ip = []

for i in range(ip_data_dict["n_sources"]):
    
    # receiver electrode locations in 2D 
    m_locs = np.vstack([
        ip_data_dict["m_locations"][i], 
        np.zeros_like(ip_data_dict["m_locations"][i])
    ]).T
    n_locs = np.vstack([
        ip_data_dict["n_locations"][i],
        np.zeros_like(ip_data_dict["n_locations"][i])
    ]).T
    
    # construct the receiver object 
    receivers = ip.receivers.Dipole(
        locations_m=m_locs, locations_n=n_locs, data_type="apparent_chargeability"
    )
    
    # construct the source 
    source = ip.sources.Dipole(
        location_a=np.r_[ip_data_dict["a_locations"][i], 0.],
        location_b=np.r_[ip_data_dict["b_locations"][i], 0.],
        receiver_list=[receivers]
    )
    
    # append the new source to the source list
    source_list_ip.append(source)

survey_ip = ip.Survey(source_list_ip)

ip_data = Data(
    survey=survey, 
    dobs=np.hstack(ip_data_dict["observed_data"]), 
    standard_deviation=np.hstack(ip_data_dict["standard_deviations"])
)

# plot psuedosection
fig, ax = plt.subplots(1, 1, figsize=(12, 4))
ip.utils.plot_pseudosection(
    ip_data, data_type="apparent chargeability",
    plot_type="contourf", data_location=True, ax=ax, 
)
ax.set_title(f"IP: {line} Pseudosection")
ax.set_aspect(1.5)  # some vertical exxageration

simulation_ip = ip.simulation_2d.Simulation2DNodal(
    mesh=mesh, solver=Solver, rho=np.exp(mopt), 
    etaMap=maps.IdentityMap(mesh), survey=survey_ip,
    storeJ=True
)

dmisfit_ip = data_misfit.L2DataMisfit(data=ip_data, simulation=simulation_ip)

# Related to inversion
reg_ip = regularization.Tikhonov(
    mesh,
    alpha_s=1/mesh.hx.min()**2.,
    alpha_x=1.,
    alpha_y=1.,
)

opt_ip = optimization.ProjectedGNCG(
    maxIter=20, maxIterCG=20, upper=np.inf, lower=0.
)

inv_prob_ip = inverse_problem.BaseInvProblem(dmisfit_ip, reg_ip, opt_ip)

m0_ip =  1e-3 * np.ones(mesh.nC)

beta_est_ip = directives.BetaEstimate_ByEig(beta0_ratio=1e0)
beta_schedule_ip = directives.BetaSchedule(coolingFactor=2, coolingRate=1)
target_ip = directives.TargetMisfit()

inv_ip = inversion.BaseInversion(
    inv_prob_ip, directiveList=[beta_est_ip, beta_schedule_ip, target_ip]
)

# Run inversion
charge_opt = inv_ip.run(m0_ip)

# plot misfit
fig, ax = plt.subplots(1, 1, figsize=(12, 4))

normalized_misfit = (inv_prob_ip.dpred - ip_data.dobs) / ip_data.standard_deviation

dc.utils.plot_pseudosection(
    ip_data, dobs=normalized_misfit, data_type="misfit",
    plot_type="contourf", data_location=True, ax=ax, 
    cbar_opts={"pad":0.25}
)

ax.set_title("IP misfit")
ax.set_aspect(1.5)  # some vertical exxageration
ax.set_xlabel("Northing (m)")

cb_axes = plt.gcf().get_axes()[-1]
cb_axes.set_xlabel('normalized misfit')

plt.tight_layout()

fig, ax = plt.subplots(1,1, figsize=(10, 5))
out = mesh.plotImage(
    charge_opt, pcolorOpts={'cmap':'Spectral_r'}, ax=ax,
    clim=(0, 30)
)
for data in geologic_section:
    ax.plot(data[:,0], data[:,1], 'k--', alpha=0.5)
ax.set_ylim((-400, 0))
ax.set_xlim(core_domain_x)
cb = plt.colorbar(out[0], label="chargeability", fraction=0.05, orientation='horizontal', ax=ax, pad=0.2)
ax.set_xlabel('Easting (m)')
ax.set_ylabel('Elevation (m)')
ax.set_aspect(1.5)
plt.tight_layout()

charge_ubc = read_ubc_model("century/46800E/IPMODA.CHG")

fig, ax = plt.subplots(2, 1, figsize=(10, 8))
rho = (mapping*mopt)
clim = (0, 30)

out = mesh.plotImage(
    charge_opt, pcolorOpts={'cmap':'Spectral_r'}, ax=ax[0],
    clim=clim
)
cb = plt.colorbar(out[0], label="Chargeability", fraction=0.05, orientation='horizontal', ax=ax[0], pad=0.2)
ax[0].set_title("SimPEG")

out = mesh_ubc.plotImage(
    charge_ubc, pcolorOpts={'cmap':'Spectral_r'}, ax=ax[1],
    clim=clim
)
cb = plt.colorbar(out[0], label="Chargeability", fraction=0.05, orientation='horizontal', ax=ax[1], pad=0.2)
ax[1].set_title("UBC")

for a in ax: 
    for data in geologic_section:
        a.plot(data[:,0], data[:,1], 'k--', alpha=0.5)
        
    a.set_xlim(core_domain_x)
    a.set_ylim((-400, 0))
    
    a.set_aspect(1.5)
    a.set_xlabel('Easting (m)')
    a.set_ylabel('Elevation (m)')
    
plt.tight_layout()