#!/usr/bin/env python3
import pandas as pd
import argparse
import re

parser = argparse.ArgumentParser()
parser.add_argument("--source", "-s", metavar=("f0(MHz)", "nuc1(i.e. 1H)"), nargs=2)
parser.add_argument("--dest", "-d", metavar="nuc2(i.e. 7Li) nucl3 ...", nargs="+")
parser.add_argument("--near", "-n", metavar="RANGE", type=float)

args = parser.parse_args()

nuclear_magneton = 7.62259328 # MHz/T
isotopes = pd.read_table('isotopetable.dat', comment="%", delim_whitespace=" ")
isotopes["gamma in MHz/T"] = isotopes["gn"] * nuclear_magneton
isotopes["stable"] = isotopes["stable"].replace(["-", "*"], ["yes","no"])
isotopes["name"] = isotopes["name"].map(lambda x : x.capitalize() )

pd.set_option("display.precision", 3)

if args.source:
    f0, nuc1 = args.source
    try:
        f0 = float(f0)
    except:
        raise ValueError("f0 is not a number")
    try:
        nuc1_n, nuc1_name = re.findall(r"(\d+)(\D+)", nuc1)[0]
        nuc1_data = isotopes.query(f'symbol == "{nuc1_name}" & nucleons == {nuc1_n}')
        if nuc1_data.empty:
            raise IndexError("Isotope not in database", nuc1)
        else:
            nuc1_gamma = nuc1_data["gamma in MHz/T"]
    except:
        raise ValueError
    # input OK calculate now B0 and the Larmor freq. of the other nuclei
    print("Source")
    isotopes["B0 in T"] = f0/float(nuc1_gamma)
    isotopes["f0 in MHz"] = isotopes["gamma in MHz/T"] * (f0/float(nuc1_gamma))
    nuc1_data = isotopes.query(f'symbol == "{nuc1_name}" & nucleons == {nuc1_n}')
    print(nuc1_data.to_markdown())

if args.dest:
    nuc_loc = []
    for nuc in args.dest:
        try:
            nuc_n, nuc_name = re.findall(r"(\d+)(\D+)", nuc)[0]
            _nuc_data = isotopes.query(f'symbol == "{nuc_name}" & nucleons == {nuc_n}')
            if _nuc_data.empty:
                # no isotope was found
                raise IndexError("Isotope not in database", nuc)
            elif _nuc_data.shape[0] > 1:
                # more than one was found (should not happen ever)
                raise IndexError("Multiple Isotopes in database, REALLY WEIRD ERROR", nuc)
            else:
                nuc_loc.append( int(_nuc_data.index.values))
                nuc_gamma = _nuc_data["gamma in MHz/T"]
        except:
            raise ValueError
    print("Destination")
    print(isotopes.loc[nuc_loc].to_markdown())

if args.near:
    near_data = isotopes.query(f'{f0 - args.near} <= `f0 in MHz` <= {f0 + args.near}')
    print(f"Nearby nuclei (+/- {args.near} MHz ")
    print(near_data.sort_values(by="f0 in MHz").to_markdown())