#!/usr/bin/python

"""
    pint-convert
    ~~~~~~~~~~~~

    :copyright: 2020 by Pint Authors, see AUTHORS for more details.
    :license: BSD, see LICENSE for more details.
"""

import argparse
import re
import sys

from pint import UnitRegistry

parser = argparse.ArgumentParser(description='Unit converter.', usage=argparse.SUPPRESS)
parser.add_argument('-s', '--system', metavar='sys', default='SI', help='unit system to convert to (default: SI)')
parser.add_argument('-p', '--prec', metavar='n', type=int, default=12, help='number of maximum significant figures (default: 12)')
parser.add_argument('-u', '--prec-unc', metavar='n', type=int, default=2, help='number of maximum uncertainty digits (default: 2)')
parser.add_argument('-U', '--no-unc', dest='unc', action='store_false', help='ignore uncertainties in constants')
parser.add_argument('-C', '--no-corr', dest='corr', action='store_false', help='ignore correlations between constants')
parser.add_argument('fr', metavar='from', type=str, help='unit or quantity to convert from')
parser.add_argument('to', type=str, nargs='?', help='unit to convert to')
try:
    args = parser.parse_args()
except SystemExit:
    parser.print_help()
    raise

ureg = UnitRegistry()
ureg.auto_reduce_dimensions = True
ureg.autoconvert_offset_to_baseunit = True
ureg.enable_contexts('Gau', 'ESU', 'sp', 'energy', 'boltzmann')
ureg.default_system = args.system

if args.unc:
    import uncertainties
    # Measured constans subject to correlation
    #  R_i: Rydberg constant
    #  g_e: Electron g factor
    #  m_u: Atomic mass constant
    #  m_e: Electron mass
    #  m_p: Proton mass
    #  m_n: Neutron mass
    R_i = (ureg._units['R_inf'].converter.scale, 0.0000000000021e7)
    g_e = (ureg._units['g_e'].converter.scale,   0.00000000000035)
    m_u = (ureg._units['m_u'].converter.scale,   0.00000000050e-27)
    m_e = (ureg._units['m_e'].converter.scale,   0.00000000028e-30)
    m_p = (ureg._units['m_p'].converter.scale,   0.00000000051e-27)
    m_n = (ureg._units['m_n'].converter.scale,   0.00000000095e-27)
    if args.corr:
        # Correlation matrix between measured constants (to be completed below)
        #          R_i       g_e       m_u       m_e       m_p       m_n
        corr = [[  1.0    , -0.00206,  0.00369,  0.00436,  0.00194,  0.00233],  # R_i
                [ -0.00206,  1.0    ,  0.99029,  0.99490,  0.97560,  0.52445],  # g_e
                [  0.00369,  0.99029,  1.0    ,  0.99536,  0.98516,  0.52959],  # m_u
                [  0.00436,  0.99490,  0.99536,  1.0    ,  0.98058,  0.52714],  # m_e
                [  0.00194,  0.97560,  0.98516,  0.98058,  1.0    ,  0.51521],  # m_p
                [  0.00233,  0.52445,  0.52959,  0.52714,  0.51521,  1.0    ]]  # m_n
        (R_i, g_e, m_u, m_e, m_p, m_n) = uncertainties.correlated_values_norm([R_i, g_e, m_u, m_e, m_p, m_n], corr)
    else:
        R_i = uncertainties.ufloat(*R_i)
        g_e = uncertainties.ufloat(*g_e)
        m_u = uncertainties.ufloat(*m_u)
        m_e = uncertainties.ufloat(*m_e)
        m_p = uncertainties.ufloat(*m_p)
        m_n = uncertainties.ufloat(*m_n)
    ureg._units['R_inf'].converter.scale = R_i
    ureg._units['g_e'].converter.scale   = g_e
    ureg._units['m_u'].converter.scale   = m_u
    ureg._units['m_e'].converter.scale   = m_e
    ureg._units['m_p'].converter.scale   = m_p
    ureg._units['m_n'].converter.scale   = m_n

    # Measured constants with zero correlation
    ureg._units['gravitational_constant'].converter.scale = uncertainties.ufloat(ureg._units['gravitational_constant'].converter.scale, 0.00015e-11)
    ureg._units['d_220'].converter.scale                  = uncertainties.ufloat(ureg._units['d_220'].converter.scale,                  0.000000032e-10)
    ureg._units['K_alpha_Cu_d_220'].converter.scale       = uncertainties.ufloat(ureg._units['K_alpha_Cu_d_220'].converter.scale,       0.00000022)
    ureg._units['K_alpha_Mo_d_220'].converter.scale       = uncertainties.ufloat(ureg._units['K_alpha_Mo_d_220'].converter.scale,       0.00000019)
    ureg._units['K_alpha_W_d_220'].converter.scale        = uncertainties.ufloat(ureg._units['K_alpha_W_d_220'].converter.scale,        0.000000098)

    ureg._root_units_cache = dict()
    ureg._build_cache()

def convert(u_from, u_to=None, unc=None, factor=None):
    q = ureg.Quantity(u_from)
    fmt = '.{}g'.format(args.prec)
    if unc:
        q = q.plus_minus(unc)
    if u_to:
        nq = q.to(u_to)
    else:
        nq = q.to_base_units()
    if (factor):
        q *= ureg.Quantity(factor)
        nq *= ureg.Quantity(factor).to_base_units()
    prec_unc = use_unc(nq.magnitude, fmt, args.prec_unc)
    if (prec_unc > 0):
        fmt = '.{}uS'.format(prec_unc)
    else:
        try:
            nq = nq.magnitude.n * nq.units
        except:
            pass
    fmt = '{:' + fmt + '} {:~P}'
    print(('{:} = ' + fmt).format(q, nq.magnitude, nq.units))

def use_unc(num, fmt, prec_unc):
    unc = 0
    try:
        if (isinstance(num, uncertainties.UFloat)):
            full = ('{:'+fmt+'}').format(num)
            unc = re.search(r'\+\/-[0.]*([\d.]*)', full).group(1)
            unc = len(unc.replace('.', ''))
    except:
        pass
    return max(0, min(prec_unc, unc))

convert(args.fr, args.to)
