"""The module provides functionality to visualise known transcripts from Ensemble database and new transcripts (the coordinates of exons should be directly provided by a user). Multiple transcripts can be plotted on one panel."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_plot_transcripts.ipynb.

# %% auto 0
__all__ = ['get_coord_from_tscrpt_id', 'draw_transcript', 'draw_exons', 'draw_transcripts_list']

# %% ../nbs/00_plot_transcripts.ipynb 3
import matplotlib.pyplot as plt
import matplotlib.lines as lines
import numpy as np
import requests, sys


# %% ../nbs/00_plot_transcripts.ipynb 5
def get_coord_from_tscrpt_id(transcipt_id):
    # Define the Ensembl REST API server and endpoint to fetch transcript details
    server = "https://rest.ensembl.org"  # Base URL for the Ensembl REST API
    ext = "/lookup/id/" + transcipt_id + "?expand=1"  # Endpoint for the transcript lookup
    # The `expand=1` query parameter ensures that additional details like exons are included in the response.

    # Make a GET request to the Ensembl API
    r = requests.get(server + ext, headers={"Content-Type": "application/json"})
    # The `headers` specify that the response should be in JSON format.

    # Check if the request was successful
    if not r.ok:  # If the HTTP status code is not OK (200),
        r.raise_for_status()  # Raise an exception with the error details
        sys.exit()  # Exit the program to prevent further execution

    # Parse the JSON response into a Python dictionary
    decoded = r.json()

    # Extract the exon data from the response
    exon_list = list(decoded['Exon'])  # Retrieve the list of exons for the transcript
    exon_coord = []  # Initialize a list to store exon coordinates

    # Iterate through each exon and extract its start and end positions
    for i, e in enumerate(exon_list):
        coord = [e.get('end'), e.get('start')]  # Extract exon coordinates as [end, start]
        exon_coord.append(coord)  # Append the coordinates to the list

    # Extract the strand information from the response
    strand = decoded['strand']  # The strand is +1 (forward) or -1 (reverse)

    # Return the list of exon coordinates and the strand information
    return (exon_coord, strand)


# %% ../nbs/00_plot_transcripts.ipynb 10
"""input:
exons: transcript's coordinates in a form of list of lists where every list represents 
coordinates of an exon (example: [[300000000, 450000000], [670000000, 1200000000], [1500000000, 1700000000], [2500000000, 3500000000], [3600000000, 3800000000]])
direction: should be 1 or -1 which corresponds to sens or antisense direction
color: any value accepted by matplotlib
"""
def draw_transcript(exons, direction, color, offset=0, no_render=False):
    # Initialize the plotting area if rendering is enabled
    if not no_render:
        plt.axes()  # Create axes for the plot
        plt.xlim((-0.1, 1))  # Set x-axis limits
        plt.ylim((-0.3, 0.3))  # Set y-axis limits
        plt.margins(0.2)  # Add margins for better spacing
        plt.axis('off')  # Hide axes
        fig = plt.gcf()  # Get the current figure
        fig.set_size_inches(20, 2)  # Set the figure size

    # Height of the rectangles representing exons
    height = 0.2
    
    # Draw a vertical line as a reference point
    plt.plot([offset + 0.1, offset + 0.1], linestyle='solid', linewidth=0.5, c='grey')
    
    # Determine indices for start and end of exons based on strand direction
    j = 0
    k = 1
    if direction == 1:  # Forward strand
        pos_start = exons[0][1]  # Start position of the first exon
        pos_end = exons[-1][0]  # End position of the last exon
    else:  # Reverse strand
        pos_start = exons[-1][0]  # Start position of the last exon
        pos_end = exons[0][1]  # End position of the first exon
        j = 1
        k = 0
    
    # Calculate total length of the transcript for scaling
    total_length = pos_end - pos_start
    total_length_with_margin = 1.05 * total_length  # Add margin for better visualization
    pos_start_with_margin = pos_start - 0.025 * total_length  # Adjust start with margin

    # Draw each exon as a rectangle
    for i, exon in enumerate(exons):
        # Calculate rectangle dimensions and position
        rectangle = plt.Rectangle(
            ((exon[j] - pos_start_with_margin) / total_length_with_margin, offset),  # Bottom-left corner
            (exon[k] - exon[j]) / total_length_with_margin,  # Width
            height,  # Height
            fc=color,  # Fill color
            ec="black"  # Edge color
        )
        plt.gca().add_patch(rectangle)  # Add the rectangle to the plot

    # Add an arrow indicating the strand direction
    if i > 0:  # Ensure there are exons to draw
        arrow = None
        if direction < 0:  # Reverse strand
            arrow = plt.arrow(1, offset - height/4, -1, 0, width=0.0015, head_length=0.01, 
                              head_width=0.1, length_includes_head=True, overhang=1)
        else:  # Forward strand
            arrow = plt.arrow(0, offset - height/4, 1, 0, width=0.0015, head_length=0.01, 
                              head_width=0.1, length_includes_head=True, overhang=1)
        plt.gca().add_patch(arrow)  # Add the arrow to the plot

    # Add vertical markers for the start and end of the transcript
    plt.plot(
        np.array([0.025, 0.025]),  # x-coordinates
        np.array([offset - height/4 - 0.03, offset - height/4 + 0.03]),  # y-coordinates
        color='black'
    )
    plt.plot(
        np.array([1 - 0.025, 1 - 0.025]),  # x-coordinates
        np.array([offset - height/4 - 0.03, offset - height/4 + 0.03]),  # y-coordinates
        color='black'
    )

    # Add labels for the start and end positions of the transcript
    plt.text(0.025, offset - height/4 - 0.075, pos_start, 
             horizontalalignment='center', verticalalignment='center', fontsize=9)
    plt.text(1 - 0.025, offset - height/4 - 0.075, pos_end, 
             horizontalalignment='center', verticalalignment='center', fontsize=9)

    # Render the plot if no_render is False
    if not no_render:
        plt.show()



# %% ../nbs/00_plot_transcripts.ipynb 14
""" 
This code provides a toolset for visualizing exons and transcripts 
# based on transcript IDs. It consists of three primary functions:
# get_coord_from_tscrpt_id(transcript_id)
"""
def get_coord_from_tscrpt_id(transcript_id):
    # Remove version number from transcript ID (e.g., 'ENST00000456328.1' -> 'ENST00000456328')
    if '.' in transcript_id:
        transcript_id = transcript_id.split('.')[0]
    
    # Define the Ensembl REST API server and endpoint to retrieve transcript details
    server = "https://rest.ensembl.org"
    ext = "/lookup/id/" + transcript_id + "?expand=1"
    
    # Send a GET request to the API with the appropriate headers
    r = requests.get(server + ext, headers={"Content-Type": "application/json"})

    # Check if the request was successful; if not, raise an error and exit
    if not r.ok:
        r.raise_for_status()
        sys.exit()

    # Parse the JSON response
    decoded = r.json()
    
    # Extract the list of exons from the response
    exon_list = list(decoded['Exon'])
    exon_coord = []

    # Iterate through each exon, extracting its start and end positions
    for i, e in enumerate(exon_list):
        coord = [e.get('end'), e.get('start')]  # Store exon coordinates as [end, start]
        exon_coord.append(coord)
    
    # Retrieve the strand direction of the transcript (+1 for forward, -1 for reverse)
    strand = decoded['strand']

    # Return a list of exon coordinates and the strand direction
    return (exon_coord, strand)


def draw_exons(exons, direction, color, transcript_name, offset=0, start_override=None, end_override=None, no_render=False):
    # Initialize a new plot if rendering is enabled
    if not no_render:
        plt.axes()
        plt.xlim((-0.1, 1))  # Set x-axis limits
        plt.ylim((-0.3, 0.3))  # Set y-axis limits
        plt.margins(0.2)  # Add plot margins
        plt.axis('off')  # Hide axes
        fig = plt.gcf()
        fig.set_size_inches(20, 2)  # Set figure size

    height = 0.2  # Height of exon rectangles

    # Draw a vertical reference line for the transcript
    plt.plot([offset + 0.1, offset + 0.1], linestyle='solid', linewidth=0.5, c='grey')

    # Determine indices for start and end coordinates based on strand direction
    j, k = 0, 1
    if direction == 1:  # Forward strand
        pos_start = exons[0][1]
        pos_end = exons[-1][0]
    else:  # Reverse strand
        pos_start = exons[-1][1]
        pos_end = exons[0][0]
        j, k = 1, 0

    # Save the real start and end positions for annotations
    real_start = pos_start
    real_end = pos_end

    # Apply start and end overrides if provided (useful for aligning multiple transcripts)
    if start_override is not None and end_override is not None:
        pos_start = start_override
        pos_end = end_override

    # Calculate scaling factors for exon positions and widths
    total_length = pos_end - pos_start
    total_length_with_margin = 1.05 * total_length
    pos_start_with_margin = pos_start - 0.025 * total_length

    # Draw each exon as a rectangle
    for i, exon in enumerate(exons):
        rectangle = plt.Rectangle(
            ((exon[j] - pos_start_with_margin) / total_length_with_margin, offset),  # Bottom-left corner
            (exon[k] - exon[j]) / total_length_with_margin,  # Width
            height,  # Height
            fc=color,  # Fill color
            ec="black"  # Edge color
        )
        plt.gca().add_patch(rectangle)  # Add rectangle to the plot

    # Draw an arrow indicating the strand direction
    if i > 0:  # Ensure there are exons to draw
        if direction < 0:  # Reverse strand
            arrow = plt.arrow(1, offset - height / 4, -1, 0, width=0.0015, head_length=0.01,
                              head_width=0.1, length_includes_head=True, overhang=1)
        else:  # Forward strand
            arrow = plt.arrow(0, offset - height / 4, 1, 0, width=0.0015, head_length=0.01,
                              head_width=0.1, length_includes_head=True, overhang=1)
        plt.gca().add_patch(arrow)

    # Annotate the start and end positions on the plot
    plt.plot(np.array([0.025 + (real_start - pos_start) / total_length / 1.05, 
                       0.025 + (real_start - pos_start) / total_length / 1.05]), 
             np.array([offset - height / 4 - 0.03, offset - height / 4 + 0.03]), color='black')
    plt.plot(np.array([1 - 0.025 - (pos_end - real_end) / total_length / 1.05, 
                       1 - 0.025 - (pos_end - real_end) / total_length / 1.05]), 
             np.array([offset - height / 4 - 0.03, offset - height / 4 + 0.03]), color='black')
    plt.text(0.025 + (real_start - pos_start) / total_length / 1.05, offset - height / 4 - 0.075, real_start,
             horizontalalignment='center', verticalalignment='center', fontsize=9)
    plt.text(1 - 0.025 - (pos_end - real_end) / total_length / 1.05, offset - height / 4 - 0.075, real_end,
             horizontalalignment='center', verticalalignment='center', fontsize=9)

    # Add transcript name annotation
    plt.text(1, offset - height, transcript_name, horizontalalignment='right', verticalalignment='top', fontsize=12)

    # Show the plot if rendering is enabled
    if not no_render:
        plt.show()


def draw_transcripts_list(transcripts_list, colors=None):
    # Prepare lists to store exon coordinates and directions for all transcripts
    exons = []
    directions = []

    # Fetch exon data and strand direction for each transcript
    for tr in transcripts_list:
        t, d = get_coord_from_tscrpt_id(tr)
        exons += [t]
        directions += [d]

    # Assign colors to transcripts if not provided
    if colors is None:
        colors = []
        for i in range(len(exons)):
            colors.append(['lightblue', 'lightgreen', 'orange', 'yellow', 'brown'][i % 5])

    # Calculate global start and end positions for alignment
    def get_limits(ex, dir):
        start = sys.maxsize
        end = -sys.maxsize
        for (e, d) in zip(ex, dir):
            if d == 1:  # Forward strand
                start = min(start, e[0][1])
                end = max(end, e[-1][0])
            else:  # Reverse strand
                start = min(start, e[-1][1])
                end = max(end, e[0][0])
        return (start, end)

    plt.axes()
    plt.xlim((-0.1, 1.1))
    plt.ylim((0.1 - 0.5 * len(exons), 0.3))
    plt.margins(0.2)
    plt.axis('off')
    fig = plt.gcf()
    fig.set_size_inches(20, len(exons) * 2)

    # Get global start and end positions
    (start, end) = get_limits(exons, directions)

    # Draw each transcript using draw_exons()
    for i, (ex, di, co, name) in enumerate(zip(exons, directions, colors, transcripts_list)):
        draw_exons(ex, di, co, name, offset=-0.5 * i, start_override=start, end_override=end, no_render=True)
        i += 1

    # Render the complete plot
    plt.show()

