# -*- coding: utf-8 -*-
"""pst_functions.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1kMRXaLXv-s42gfy4BTwjVGjzQ6sxnz2d

Install the libraries, do we need to install libaries here, no I geuss?

## Function to create an IGV Plot using Pf7

### Setup
"""

import malariagen_data
# Install igv notebook using the python command 'pip'
import igv_notebook

def sample_viewer(sample_id, chromosome, start_base_coordinate, end_base_coordinate):
  """
    Visualize genomic data using IGV_notebook.

    Parameters:
    - sample_id (str): Sample ID given in Pf7 for the sample.
    - chromosome (int): The chromosome number.
    - start_base_coordinate (int): The starting base coordinate for visualization.
    - end_base_coordinate (int): The ending base coordinate for visualization.

    Returns:
    - igv_plot: IGV Browser object for further interaction.
  """
  # Access to annotation file
  annotation_path = '/content/Pfalciparum_replace_Pf3D7_MIT_v3_with_Pf_M76611.gff'

  # Access to reference file
  reference_url = 'https://storage.googleapis.com/pf7_release/reference/Pfalciparum.genome.fasta'

  # Define the genomic locus formatted for igv-notebook
  locus = f'Pf3D7_{str(chromosome).zfill(2)}_v3:{start_base_coordinate:,}-{end_base_coordinate:,}'
  print(f'Visualizing chromosome {chromosome} for genomic locus: {locus}')

  # Initialize the browser
  igv_notebook.init()
  igv_plot = igv_notebook.Browser(

    {
        "locus": locus,
            "tracks": [

            {
                "name": "Annotations",
                "path": annotation_path,
                "indexed": False,
                "type": "annotation",
                "format": "gff"
            },
            {
                "name": "BAM",
                "url": f"https://pf7_release.cog.sanger.ac.uk/bam/{sample_id}.bam",
                "indexURL": f"https://pf7_release.cog.sanger.ac.uk/bam/{sample_id}.bam.bai",
                "format": "bam",
                "type": "alignment"
            },
        ],

        # Load the reference
        "reference": {
           "fastaURL": reference_url,
           "indexed": False, # reads entire fasta
        },
    })
  return igv_plot

"""## Functions for Mapping Drug Resistance Frequency Across West Africa

### Setup
"""

import collections
import pandas as pd
import numpy as np

drugs = [
    'Chloroquine',
    'Pyrimethamine',
    'Sulfadoxine',
    'Mefloquine',
    'Artemisinin',
    'Piperaquine',
    'SP (uncomplicated)',
    'SP (IPTp)',
]

"""Here we define two functions **proportion_agg** and **n_agg**.

Proportion_agg summarises the proportion of samples listed as 'Resistant' out of all samples which are listed as either 'Sensitive' or 'Resistant' - i.e. not 'Undetermined'. It returns a pandas series object with this information.

n_agg does the same but for the counts, rather than proportions.
"""

def proportion_agg(x):
    names = collections.OrderedDict() # create an empty ordered dictionary
    for drug in drugs: # Loop over each drug type
        n = np.count_nonzero( (x[drug] != 'Undetermined') ) # Count how many entries are not 'Undetermined'
        if n == 0:
            proportion = np.nan # Assign nan if none
        else:
            proportion = np.count_nonzero( # Otherwise, calculate the proportion of samples which are 'Resistant'
                ( x[drug] == 'Resistant')
            ) / np.count_nonzero(
                ( x[drug] != 'Undetermined' )
            )

        names[drug] = proportion
    return pd.Series(names)

def n_agg(x):
    names = collections.OrderedDict()
    for drug in drugs:
        n = np.count_nonzero( (x[drug] != 'Undetermined') )
        names[drug] = n
    return pd.Series(names)

"""Plot Pie Setup"""

#import plotting libraries
import matplotlib.pyplot as plt
import cartopy
import geopandas

#specifc imports to map and plot samples
import pandas as pd
import numpy as np
import collections
import malariagen_data
import warnings
import matplotlib.pyplot as plt
import cartopy
import geopandas
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from matplotlib.offsetbox import AnchoredText
from collections import Counter
from cartopy.io import shapereader
from time import strftime
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
import matplotlib.patches as mpatches
from matplotlib.image import imread

# Set the figure resolution - higher values = higher resolution
plt.rcParams['figure.dpi'] = 200
fig = plt.figure()

# Create a list of dictionaries 'location_coords'
# This list stores information for each location in the plot:
# annotation_ha sets the alignment of the pie chart with a location
# x/y_offsets set how far off the centre alignment the pie chart will appear
# pie_lon, pie_lat set the coordinates of each pie chart
# connection_style sets the connector line

# Note that here the administrative divisions are hard coded - i.e. written out rather than linked to a variable
# If you were to investigate a different region, you may need to re-write these names according to
# what is contiained within the 'locations' variable, as well as tweaking the other values
# to suit your new map

locations_coords = [
        dict(location = ('Dakar'), annotation_ha='center' ,         x_offset=-2 , y_offset=1.5, pie_lon=-0.7071, pie_lat=0.7071, connection_style = "arc3,rad=0."),
        dict(location = ('Western'), annotation_ha='center' ,       x_offset=-2 , y_offset=-1.5, pie_lon=-1, pie_lat=0, connection_style ="arc3,rad=0."),
        dict(location = ('Sedhiou'), annotation_ha='center' ,       x_offset=-0.5 , y_offset=-2.2, pie_lon=-0.7071, pie_lat=-0.7071, connection_style ="arc3,rad=0."),
        dict(location = ('Upper River'), annotation_ha='center',    x_offset=0 , y_offset=2.5, pie_lon=0, pie_lat=1, connection_style ="arc3,rad=0."),
        dict(location = ('Kayes') , annotation_ha='center',         x_offset=-1.5 , y_offset=-2.0, pie_lon=-0.7071, pie_lat=-0.7071, connection_style ="arc3,rad=-0."),
        dict(location = ('Hodh el Gharbi'), annotation_ha='center', x_offset=3.2 , y_offset=0, pie_lon=1, pie_lat=0, connection_style ="arc3,rad=-0."),
        dict(location = ('Bamako'), annotation_ha='center' ,        x_offset=-1.5 , y_offset=-2, pie_lon=-0.7071, pie_lat=-0.7071, connection_style ="arc3,rad=-0."),
        dict(location = ('Koulikoro'), annotation_ha='center' ,     x_offset=2.5 , y_offset=0.8, pie_lon=0.7071, pie_lat=0.7071, connection_style ="arc3,rad=-0."),
        dict(location = ('Sikasso'),annotation_ha='center' ,        x_offset=0 , y_offset=-2.5, pie_lon=0, pie_lat=-1, connection_style ="arc3,rad=0."),
        dict(location = ('Abidjan'), annotation_ha='center' ,       x_offset=-3 , y_offset=0, pie_lon=-1, pie_lat=0, connection_style ="arc3,rad=0."),
        dict(location = ('Ashanti'), annotation_ha='center' ,       x_offset=-2.2 , y_offset=1, pie_lon=-0.7071, pie_lat=0.7071, connection_style ="arc3,rad=0."),
        dict(location = ('Central'),annotation_ha='center' ,        x_offset=-2 , y_offset=-2, pie_lon=-0.7071, pie_lat=-0.7071, connection_style ="arc3,rad=0."),
        dict(location = ('Upper East'), annotation_ha='center' ,    x_offset=0 , y_offset=2.5, pie_lon=0, pie_lat=1, connection_style ="arc3,rad=0."),
        dict(location = ('Greater Accra'), annotation_ha='center' , x_offset=0 , y_offset=-2.5, pie_lon=0, pie_lat=-1, connection_style ="arc3,rad=0."),
        dict(location = ('Volta'),annotation_ha='center' ,          x_offset=0 , y_offset=2.3, pie_lon=0, pie_lat=1, connection_style ="arc3,rad=0."),
        dict(location = ('Atlantique'), annotation_ha='center' ,    x_offset=2 , y_offset=2, pie_lon=0.7071, pie_lat=0.7071, connection_style ="arc3,rad=-0."),
        dict(location = ('Littoral'), annotation_ha='center' ,      x_offset=2 , y_offset=-2, pie_lon=0.7071, pie_lat=-0.7071, connection_style ="arc3,rad=-0."),
        dict(location = ('Sud-Ouest'),annotation_ha='center' ,      x_offset=0 , y_offset=2.5, pie_lon=0, pie_lat=1, connection_style ="arc3,rad=0."),
        dict(location = ('Wouleu-Ntem'),annotation_ha='center' ,    x_offset=-2 , y_offset=1, pie_lon=0, pie_lat=1, connection_style ="arc3,rad=-0."),
]

def plot_pie(proportion_resistant,lon,lat,ax,width,height,n, loc):
    adm1 = loc.split(',')[0]
    ax_sub = inset_axes(
        ax, width, height, loc=10,
        bbox_to_anchor=(lon + locations_coords[n]['pie_lon'], lat + locations_coords[n]['pie_lat']),
        bbox_transform=ax.transData,
        borderpad=0
    )
    # Set colours for the pie chart
    wedges,texts= ax_sub.pie(proportion_resistant, colors = ['red', 'grey'], radius=2.75, startangle=90, counterclock=False)

    # Set the line connector between pie chart and location
    ax.plot(
        [lon, lon + locations_coords[n]['pie_lon']]
        , [lat, lat + locations_coords[n]['pie_lat']]
        , color='black'
    )
    # Set annotations for the pie chart
    ax.annotate(loc.split(', ')[0] + ',\n' + loc.split(', ')[1].split('(')[0] + '\n(' + loc.split('(')[1],
                xy=(lon, lat),
                xycoords='data',
                xytext=(
                    lon + locations_coords[n]['x_offset'],
                    lat + locations_coords[n]['y_offset']
                ),
                textcoords='data',
                fontsize=4,
                ha = locations_coords['location' == adm1]['annotation_ha'],
                va = 'center'
                )

def plot_chloroquine_resistance_map(image, resistance_proportions, locations, df_all_sample_metadata):
    # Set up the plot axes using a Plate Carrée projection
    ax = plt.axes(projection=ccrs.PlateCarree())
    ax.imshow(imread(image), origin='upper', transform=ccrs.PlateCarree(), extent=[-180, 180, -90, 90])
    ax.add_feature(cartopy.feature.BORDERS, linestyle=':', linewidth=0.6, alpha=0.4)

    # Make a new dataframe 'proportion_resistant' which sets a new index and drops the old index
    proportion_resistant = resistance_proportions.reset_index(drop=True)

    # Iterate over locations and plot pie
    for n, loc in enumerate(locations):
        adm1 = loc.split(',')[0]  # Save the part of the 'loc' string occurring before the comma ',' as adm1

        # Extract latitude and longitude for the location
        lat = np.unique(df_all_sample_metadata.loc[df_all_sample_metadata['Admin level 1'] == adm1]['Admin level 1 latitude'])
        lon = np.unique(df_all_sample_metadata.loc[df_all_sample_metadata['Admin level 1'] == adm1]['Admin level 1 longitude'])

        # Define the function 'plot_pie' to make a pie chart per location
        # This function takes parameters including the proportion of resistant samples,
        # The location coordinates (lon, lat)
        # ax = location on the larger map of West Africa
        # width, height = size of pie chart
        plot_pie([proportion_resistant[n], 1 - proportion_resistant[n]], lon, lat, ax, 0.08, 0.08, n, loc)

    # Set final plot axes limits
    ax.set_xlim(-21, 13)
    ax.set_ylim(0, 18)

    # Set a legend
    red_patch = mpatches.Patch(color='red', label='Chloroquine resistant')
    plt.legend(handles=[red_patch], loc=(-55, -3), prop={'size': 6})
    plt.show()

"""## Functions for Species Detection"""

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.lines import Line2D
from time import sleep

def quick_align(seqA, seqB):
    return "".join([("|" if pair[0] == pair[1] else
                     (" " if " " in pair else "*")) for pair in zip(seqA, seqB)])

def rc(s):
    d = { "A": "T", "T": "A", "C": "G", "G": "C"}
    return "".join([d[l] for l in s[::-1] if l in d.keys()])

def display_alignment(s1, s2):
    print("-" * 50)
    print("Displaying alignment...")
    print("")
    print(s1)
    print(quick_align(s1, s2))
    print(s2)
    print("-" * 50)
    print("")
    print("")


def colour_coded_alignment(s1, s2):
    alignment = quick_align(s1, s2)

    if (len(s1) == len(s2) == len(alignment)) == False:
        print("-" * 50)
        print("WARNING - colour_code_alignment function was not executed!")
        print("-" * 50)
        print("Try to add spaces to either Reference or Sequence1 above so that:")
        print(" - they contain the same number of characters (spaces included)")
        print(" - and you have minimised the number of mismatches")
        print("-" * 50)

        return

    xs = np.arange(len(s1))

    alt_A = [(p[2] == "*") & (p[1] == "A") for p in zip(s1, s2, alignment)]
    alt_T = [(p[2] == "*") & (p[1] == "T") for p in zip(s1, s2, alignment)]
    alt_C = [(p[2] == "*") & (p[1] == "C") for p in zip(s1, s2, alignment)]
    alt_G = [(p[2] == "*") & (p[1] == "G") for p in zip(s1, s2, alignment)]
    deletion = [p[1] == " " for p in zip(s1, s2, alignment)]
    insertion = [p[0] == " " for p in zip(s1, s2, alignment)]
    ref = [p[2] == "|" for p in zip(s1, s2, alignment)]

    fig, ax = plt.subplots(figsize = (5, 1))

    ax.scatter(xs[alt_A], np.zeros(xs[alt_A].size), color = "lime", marker = "s")
    ax.scatter(xs[alt_T], np.zeros(xs[alt_T].size), color = "red", marker = "s")
    ax.scatter(xs[alt_C], np.zeros(xs[alt_C].size), color = "blue", marker = "s")
    ax.scatter(xs[alt_G], np.zeros(xs[alt_G].size), color = "goldenrod", marker = "s")
    ax.scatter(xs[deletion], np.zeros(xs[deletion].size), color = "black", marker = "s")
    ax.scatter(xs[insertion], np.zeros(xs[insertion].size), color = "mediumpurple", marker = "s")
    ax.scatter(xs[ref], np.zeros(xs[ref].size), color = "lightgrey", zorder = -1, marker = "s")

    ax.set_yticks([])
    ax.axis("off")


    print("-" * 50)
    print("Displaying colour-coded alignment...")
    print("-" * 50)
    plt.show()
    print("-" * 50)

    return

def show_species_alignments_to_reference():
    Pv = "GGGAATAATCTGTGTTACAATACTTTTTGATCCCAGGCTGGTAAAAAATGTAAACTTTTAGCCCAAAAGAATAG"

    Po = "GGGAATTATCTGTGTTACAATACTTTTTGATCCCAGGCTGGTAAAAAATGTAAACTTTTAGCCCAATAGAATAG"
    Pm = "GGGAATTATCTGTGTTACAATA TTTTTGATCCCAGGCTGGTATTAAATGTAAACTTTTAGCCCAAAAGAATAG"
    Pk = "GGGAATAATCTGTGTTACAATACTTTTTGATCCCAGGCTGGTAAAAAATGTAAACTTTTAGCCCGAAAGAATAG"
    Pf = "GGGAATTATCTGTGTTACAA A TTTTTGATCCCAGGCTGGTAAAAAATGTAAACTTTTAGCCCATAAGAATAG"

    species_dict = {"Pf": "Plasmodium falciparum",
                    "Pv": "Plasmodium vivax",
                    "Po": "Plasmodium ovale",
                    "Pm": "Plasmodium malariae",
                    "Pk": "Plasmodium knowlesi"}

    legend_dict = {"Substituted to A": "lime",
                   "Substituted to T": "red",
                   "Substituted to C": "blue",
                   "Substituted to G": "goldenrod",
                   "Deletion": "black",
                   "Insertion": "mediumpurple"}

    legend_items = [Line2D([0], [0], color=v, lw=4, label=k) for k, v in legend_dict.items()]

    fig, ax = plt.subplots(6, 1, figsize = (5, 6))

    ax[0].legend(handles=legend_items, loc='center', ncol = 3)
    ax[0].set_yticks([])
    ax[0].axis("off")

    for i, species in enumerate(["Pv", "Po", "Pk", "Pm", "Pf"], 1):

        s1 = Pv
        s2 = eval(species)
        xs = np.arange(len(s1))

        alignment = quick_align(s1, s2)

        alt_A = [(p[2] == "*") & (p[1] == "A") for p in zip(s1, s2, alignment)]
        alt_T = [(p[2] == "*") & (p[1] == "T") for p in zip(s1, s2, alignment)]
        alt_C = [(p[2] == "*") & (p[1] == "C") for p in zip(s1, s2, alignment)]
        alt_G = [(p[2] == "*") & (p[1] == "G") for p in zip(s1, s2, alignment)]
        deletion = [p[1] == " " for p in zip(s1, s2, alignment)]
        insertion = [p[0] == " " for p in zip(s1, s2, alignment)]
        ref = [p[2] == "|" for p in zip(s1, s2, alignment)]

        ax[i].scatter(xs[alt_A], np.zeros(xs[alt_A].size), color = "lime", marker = "s")
        ax[i].scatter(xs[alt_T], np.zeros(xs[alt_T].size), color = "red", marker = "s")
        ax[i].scatter(xs[alt_C], np.zeros(xs[alt_C].size), color = "blue", marker = "s")
        ax[i].scatter(xs[alt_G], np.zeros(xs[alt_G].size), color = "goldenrod", marker = "s")
        ax[i].scatter(xs[deletion], np.zeros(xs[deletion].size), color = "black", marker = "s")
        ax[i].scatter(xs[insertion], np.zeros(xs[insertion].size), color = "mediumpurple", marker = "s")
        ax[i].scatter(xs[ref], np.zeros(xs[ref].size), color = "lightgrey", zorder = -1, marker = "s")

        ax[i].text(0, 0.02, s = species_dict[species])

        ax[i].set_yticks([])
        ax[i].axis("off")
    plt.show()

    return

def reveal_answer():
    print("-" * 50)
    sleep(1)
    print("The answer is")
    sleep(1)
    print("...")
    sleep(1)
    print("...")
    sleep(1)
    print("...")
    sleep(1)
    print("Plasmodium falciparum!")
    sleep(1)
    print("-" * 50)

    return