from math import pi
from io import BytesIO
import base64
from django.shortcuts import render
from django.utils.safestring import mark_safe
import re
import numpy as np
import matplotlib
from setuptools.command.rotate import rotate

matplotlib.use('Agg')
from matplotlib.patches import Wedge, FancyArrowPatch
import matplotlib.pyplot as plt

from isotopapp.models import Isotope
# Create your views here.

def home(request):
    isotopes = [i for i in Isotope.objects.all() if (i.gamma != 0 or i.stable)]
    return render(request, 'home.html', {'isotopes': [[f"{i.n_nucleons}{i.symbol}", mark_safe(f"<sup>{i.n_nucleons}</sup>{i.symbol}")] for i in isotopes]})

def extract_isotope_parts(isotope_str):
    """Extracts the number and element from an isotope string (e.g., '23Na')."""
    print(isotope_str)
    match = re.match(r"(\d+)([A-Za-z]+)", isotope_str)
    if not match:
        raise ValueError("Invalid isotope format")
    return int(match.group(1)), match.group(2)

def isotope_info(isotope, field):
    f_larmor = field*isotope.gamma
    if isotope.spin_quantum_number.as_integer_ratio()[1]==1:
        spin = int(isotope.spin_quantum_number)
    else:
        spin = mark_safe(f"<sup>{isotope.spin_quantum_number.as_integer_ratio()[0]}</sup>&frasl;<sub>{isotope.spin_quantum_number.as_integer_ratio()[1]}</sub>")
    return [isotope.n_nucleons,
            mark_safe(f"<sup>{isotope.n_nucleons}</sup>{isotope.symbol}"),
            isotope.name.capitalize(),
            f"{f_larmor:.3f}",
            spin,
            f"{isotope.natural_abundance:.1f}",
            f"{isotope.gamma:.5f}"]

def result(request):
    n1, element1 = extract_isotope_parts(request.GET.get('isotope1'))
    isotope1 = Isotope.objects.filter(symbol=element1, n_nucleons=n1).get()

    freq = float(request.GET.get('freq'))

    field_T = freq / isotope1.gamma

    figure=None
    print(request.GET)
    if request.GET.get('range_search') == "":
        close_isotopes = [isotope_info(isotope1, field_T)]
        freq_range = float(request.GET.get('freq_range'))
        Isotope.objects.filter()
        # calculate the frequency for all isotopes and compile a list of close by isotopes
        for isotope in Isotope.objects.all():
            if isotope.gamma == 0: continue
            if not isotope.stable: continue
            f_Larmor = field_T*isotope.gamma
            if abs(f_Larmor - freq) <= freq_range:
                close_isotopes.append(isotope_info(isotope, field_T))
        ans = sorted(close_isotopes, key=lambda x: -float(x[3]))

    elif request.GET.get('gradient_search') == "":
        fig, ax = plt.subplots(1, 1)
        theta1, theta2 = 0, 360
        radius = 2.5
        center = (0, 0)
        w = Wedge(center, radius, theta1, theta2, fc='#D5D9FFB2', edgecolor='black')
        w_upper = Wedge((0,2.5), radius, theta1, theta2, fc='#D5D9FF42', edgecolor='black')
        w_lower = Wedge((0,-2.5), radius, theta1, theta2, fc='#D5D9FF42', edgecolor='black')

        ax.set_xlim([-8, 8])
        ax.set_ylim([-8, 8])
        ax.grid(True)

        ax.add_patch(w_upper)
        ax.add_patch(w_lower)
        ax.add_patch(w)

        sample_diameter = 5e-3
        gradient = float(request.GET.get('gradient'))

        close_isotopes = []
        for isotope in Isotope.objects.all():
            if isotope.gamma == 0: continue
            if not isotope.stable: continue
            z = (freq/isotope.gamma-field_T)/gradient
            if abs(z) <= sample_diameter:
                i_info = isotope_info(isotope, field_T)
                #i_info[3] = f"{z*1e3:.1f} mm"
                close_isotopes.append(i_info)

                arr = FancyArrowPatch((-2.5,z*1e3), (2.5, z*1e3),
                               arrowstyle='<->,head_width=.1', mutation_scale=20, linewidth=2)
                ax.add_patch(arr)
                ax.annotate(f"{isotope.n_nucleons}{isotope.symbol}", (1.1, 0), xycoords=arr, ha='left', va='center',fontsize=8)

        ax.set_aspect('equal')

        ax.set_ylabel('z (mm)')

        buf = BytesIO()
        plt.savefig(buf, format='png')
        figure = base64.b64encode(buf.getvalue()).decode('utf-8').replace('\n', '')
        buf.close()
        ans = sorted(close_isotopes, key=lambda x: float(x[3]))

    elif request.GET.get('transform') == "":
        n2, element2 = extract_isotope_parts(request.GET.get('isotope2'))
        isotope2 = Isotope.objects.filter(symbol=element2, n_nucleons=n2).get()
        #isotope_info(isotope2, field_T)
        ans = [isotope_info(isotope2, field_T)]
    else:
        ans = []
    return render(request, 'result.html', {'ans': ans, 'figure': figure})