#!/usr/bin/env python
# coding: utf-8

# # Airborne magnetic survey of Britain
# 
# This is a digitized version of an airborne magnetic survey of Britain. The exact date of measurements is not available (only the year). Flight height is estimated and not specified if it's geometric or orthometric. Given the error involved in how the data were digitized, it likely doesn't matter which. The horizontal datum is OSGB36 (epsg:27700) and we convert the data to WGS84 for standardization.
# 
# Contains British Geological Survey materials © UKRI 2021
# 
# License: [Open Government Licence](https://www.bgs.ac.uk/bgs-intellectual-property-rights/open-government-licence/)

# In[1]:


import os
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import verde as vd
import pooch
import pyproj
import pygmt


# ## Download the data archive
# 
# Use Pooch to download the data archive and fetch the CSV file only (not the PDF reports included).

# In[2]:


fnames = pooch.retrieve(
    url="https://www.bgs.ac.uk/?wpdmdl=11839",
    fname="britain-aeromag.zip",
    known_hash="md5:e9968928e1c156bfe318bfb84657a472",
    # Unpack and return only the actual data, not the reports.
    processor=pooch.Unzip(members=["Aeromag_csv/aeromag.csv"]),
)
fname = fnames[0]


# In[3]:


print(f"size: {os.path.getsize(fname) / 1e6} Mb")


# ## Read the data
# 
# Use pandas to load the CSV file.

# In[4]:


data = pd.read_csv(
    fname, 
    usecols=[0, 1, 4, 5, 10, 13], 
    header=0,
    names=["survey", "line_and_segment", "longitude_osgb", "latitude_osgb", "height_m", "total_field_anomaly_nt"],
)
data


# ## Extract the year information from the survey area field
# 
# The year is encoded in the `survey` variable. Split it up and make it a separate field.

# In[5]:


data["year"] = data.survey.map(lambda x: int(x.split("_")[0][2:]) + 1900)
data


# ## Transform to WGS84
# 
# Transform the horizontal datum from OSGB36 to WGS84. See [this StackOverflow answer for the transformation](https://stackoverflow.com/a/1434973).

# In[6]:


osgb = pyproj.CRS.from_proj4("+proj=latlong +ellps=airy +no_defs +towgs84=446.448,-125.157,542.060,0.1502,0.2470,0.8421,-20.4894")
wgs84 = pyproj.CRS.from_epsg(4326)
transformer = pyproj.Transformer.from_crs(osgb, wgs84, always_xy=True)


# In[7]:


data["longitude"], data["latitude"] = transformer.transform(data.longitude_osgb, data.latitude_osgb)


# In[8]:


data


# The difference is very small but we may as well do it.

# ## Plot the data
# 
# Make a quick plot to make sure the data look OK.

# In[9]:


fig = pygmt.Figure()
scale = np.percentile(data.total_field_anomaly_nt, 95)
pygmt.makecpt(cmap="polar", series=[-scale, scale])
fig.plot(x=data.longitude, y=data.latitude, style="c0.02c", color=data.total_field_anomaly_nt, cmap=True, projection="M15c")
fig.colorbar(frame='af+l"nT"', position="jBL+h+w7c/0.2c+o1/2")
fig.coast(shorelines=True)
fig.basemap(frame="afg")
fig.savefig("preview.jpg", dpi=200)
fig.show()


# ## Export
# 
# Make a separate DataFrame to export to CSV with only the chosen fields.

# In[10]:


export = pd.DataFrame({  
    "line_and_segment": data.line_and_segment, 
    "year": data.year.map(lambda x: "{:d}".format(x)),  
    "longitude": data.longitude.map(lambda x: "{:.5f}".format(x)),
    "latitude": data.latitude.map(lambda x: "{:.5f}".format(x)),    
    "height_m": data.height_m.map(lambda x: "{:.0f}".format(x)),    
    "total_field_anomaly_nt": data.total_field_anomaly_nt.map(lambda x: "{:.0f}".format(x)), 
})
export


# In[11]:


fname = "britain-magnetic.csv.xz"
export.to_csv(fname, index=False)
print(fname)
print(f"size: {os.path.getsize(fname) / 1e6} Mb")
for alg in ["md5", "sha256"]:
    print(f"{alg}:{pooch.file_hash(fname, alg=alg)}")


# ## Read back the data and plot it
# 
# Verify that the output didn't corrupt anything.

# In[12]:


mag = pd.read_csv(fname)
mag


# In[13]:


fig = pygmt.Figure()
scale = np.percentile(mag.total_field_anomaly_nt, 95)
pygmt.makecpt(cmap="polar", series=[-scale, scale])
fig.plot(x=mag.longitude, y=mag.latitude, style="c0.02c", color=mag.total_field_anomaly_nt, cmap=True, projection="M15c")
fig.colorbar(frame='af+l"nT"', position="jBL+h+w7c/0.2c+o1/2")
fig.coast(shorelines=True)
fig.basemap(frame="afg")
fig.show()