"""Units of measure module."""

from pathlib import Path
import json
import warnings
from functools import lru_cache
from resqpy.olio.exceptions import InvalidUnitError, IncompatibleUnitsError

version = '6th July 2021'

# physical constants
feet_to_metres = 0.3048
metres_to_feet = 1.0 / feet_to_metres
ft3_to_m3 = 0.028316846592
m3_to_ft3 = 1.0 / ft3_to_m3
bbl_to_m3 = 0.158987294928
m3_to_bbl = 1.0 / bbl_to_m3
psi_to_kPa = 44.482216152605 / 6.4516
kPa_to_psi = 1.0 / psi_to_kPa
d_to_s = float(24 * 60 * 60)
s_to_d = 1.0 / d_to_s

# Mapping from uom to set of common case-insensitive aliases
# Nb. No need to write out fractional combinations such as "bbl/day"
UOM_ALIASES = {
   # Lengths
   'm': {'m', 'metre', 'metres', 'meter', 'meters'},
   'ft': {'ft', 'foot', 'feet'},

   # Times
   'ms': {'ms', 'msec', 'millisecs', 'millisecond', 'milliseconds'},
   's': {'s', 'sec', 'secs', 'second', 'seconds'},
   'min': {'min', 'mins', 'minute', 'minutes'},
   'h': {'h', 'hr', 'hour', 'hours'},
   'd': {'day', 'days'},
   'wk': {'wk', 'week', 'weeks'},
   'a': {'a', 'yr', 'year', 'years'},

   # Ratios
   '%': {'%', 'pu', 'p.u.'},
   'm3/m3': {'m3/m3', 'v/v'},
   'g/cm3': {'g/cm3', 'g/cc'},

   # Volumes
   'bbl': {'bbl', 'stb'},
   '1000 bbl': {'1000 bbl', 'mstb', 'mbbl'},
   '1E6 bbl': {'1E6 bbl', 'mmstb', 'mmbbl'},
   '1E6 ft3': {'1E6 ft3', 'mmscf'},
   '1000 ft3': {'1000 ft3', 'mscf'},
   'm3': {'m3', 'sm3'},
   'ft3': {'ft3', 'scf'},

   # Other
   'gAPI': {'gapi'},
   'S': {'mho'},
   'mS': {'mmho'},
   'psi': {'psi', 'psia'},
   'Euc': {'count'},
}
# Mapping from alias to valid uom
UOM_ALIAS_MAP = {alias.casefold(): uom for uom, aliases in UOM_ALIASES.items() for alias in aliases}

# Set of uoms that can be safely matched case-insensitive
CASE_INSENSITIVE_UOMS = {'m', 'ft', 'm3', 'ft3', 'm3/m3', 'ft3/ft3', 'bbl', 'bar', 'psi', 'm3/d', 'bbl/d'}


@lru_cache(None)
def rq_uom(units, quantity = None):
   """Returns RESQML uom string equivalent to units
   
   Args:
      units (str): unit to coerce
      quantity (str, optional): if given, raise an exception if the uom is not supported
         for this quantity

   Returns:
      str: unit of measure

   Raises:
      InvalidUnitError: if units cannot be coerced into RESQML units for the given quantity
   """
   if not units:
      raise InvalidUnitError("Must provide non-empty unit")

   uom = _try_parse_unit(units)

   # May be a fraction: match each part against known aliases
   if uom is None and '/' in units:
      parts = units.split('/', 1)
      newpart0 = _try_parse_unit(parts[0])
      newpart1 = _try_parse_unit(parts[1])
      uom = _try_parse_unit(f"{newpart0}/{newpart1}")

   if uom is None:
      raise InvalidUnitError(f"Cannot coerce {units} into a valid RESQML unit of measure.")

   if quantity is not None:
      supported_uoms = valid_uoms(quantity = quantity)
      if uom not in supported_uoms:
         raise InvalidUnitError(f"Unit {uom} is not supported for quantity {quantity}.\n"
                                f"Supported units:\n{supported_uoms}")
   return uom


def convert(x, unit_from, unit_to, quantity = None, inplace = False):
   """Convert value between two compatible units

   Args:
      x (numeric or np.array): value(s) to convert
      unit_from (str): resqml uom
      unit_to (str): resqml uom
      quantity (str, optional): If provided, raise an exception if units are not supported
         by this quantity
      inplace (bool): if True, convert arrays in-place. Else, return new value

   Returns:
      Converted value(s)

   Raises:
      InvalidUnitError: if units cannot be coerced into RESQML units
      IncompatibleUnitsError: if units do not have compatible base units
   """

   # conversion data assume the formula "y=(A + Bx)/(C + Dx)" where "y" represents a value in the base unit.
   # Backwards formula: x=(A-Cy)/(Dy-B)
   # All current units have D==0

   uom1 = rq_uom(unit_from, quantity = quantity)
   uom2 = rq_uom(unit_to, quantity = quantity)

   if uom1 == uom2:
      return x

   base1, dim1, (A1, B1, C1, D1) = get_conversion_factors(uom1)
   base2, dim2, (A2, B2, C2, D2) = get_conversion_factors(uom2)

   if base1 != base2:
      if dim1 != dim2:
         raise IncompatibleUnitsError(f"Cannot convert from '{unit_from}' to '{unit_to}':"
                                      f"\n - '{uom1}' has base unit '{base1} and dimension '{dim1}'."
                                      f"\n - '{uom2}' has base unit '{base2} and dimension '{dim2}'.")
      else:
         warnings.warn(f"Assuming base units {base1} and {base2} are equivalent as they have the same dimensions:"
                       f"\n - '{uom1}' has base unit '{base1} and dimension '{dim1}'."
                       f"\n - '{uom2}' has base unit '{base2} and dimension '{dim2}'.")

   if not inplace:
      y = (A1 + (B1 * x)) / (C1 + (D1 * x))
      return (A2 - (C2 * y)) / ((D2 * y) - B2)

   else:
      if any(f != 0 for f in [A1, A2, D1, D2]):
         raise NotImplementedError("In-place conversion not yet implemented for non-trivial conversions")

      factor = (B1 * C2) / (C1 * B2)
      x *= factor
      return x


@lru_cache(None)
def valid_uoms(quantity = None, return_attributes = False):
   """Return set of valid RESQML units of measure
   
   Args:
      quantity (str): If given, filter to uoms supported by this quanitity.
      return_attributes (bool): If True, return a dict of all uoms and their
         attributes, such as the full name and dimension. Else, simply return
         the set of valid uoms.

   Returns
      set or dict
   """
   uoms = _properties_data()['units']
   if quantity:
      all_quantities = _properties_data()['quantities']
      supported_members = all_quantities[quantity]['members']
      uoms = {k: v for k, v in uoms.items() if k in supported_members}
   if return_attributes:
      return uoms
   else:
      return set(uoms.keys())


@lru_cache(None)
def valid_quantities(return_attributes = False):
   """Return set of valid RESQML quantities
   
   Args:
      return_attributes (bool): If True, return a dict of all quantities and their
         attributes, such as the supported units of measure. Else, simply return
         the set of valid properties.
   
   Returns
      set or dict
   """
   quantities = _properties_data()['quantities']
   if return_attributes:
      return quantities
   else:
      return set(quantities.keys())


def valid_property_kinds():
   """Return set of valid property kinds"""

   return set(_properties_data()['property_kinds'].keys())


def rq_uom_list(units_list):
   """Returns a list of RESQML uom equivalents for units in list."""

   return [rq_uom(u) for u in units_list]


def rq_length_unit(units):
   """Returns length units string as expected by resqml."""

   return rq_uom(units, quantity = 'length')


def rq_time_unit(units):
   """Returns time units string as expected by resqml."""

   return rq_uom(units, quantity = 'time')


def convert_times(a, from_units, to_units, invert = False):
   """Converts values in numpy array (or a scalar) from one time unit to another, in situ if array.
   
   note:
      To see supported units, use: `valid_uoms(quantity='time')`
   """

   if invert:
      from_units, to_units = to_units, from_units

   return convert(a, from_units, to_units, quantity = 'time', inplace = True)


def convert_lengths(a, from_units, to_units):
   """Converts values in numpy array (or a scalar) from one length unit to another, in situ if array.

   arguments:
      a (numpy float array, or float): array of length values to undergo unit conversion in situ, or a scalar
      from_units (string): the units of the data before conversion
      to_units (string): the required units

   returns:
      a after unit conversion

   note:
      To see supported units, use: `valid_uoms(quantity='length')`
   """

   return convert(a, from_units, to_units, quantity = 'length', inplace = True)


def convert_pressures(a, from_units, to_units):
   """Converts values in numpy array (or a scalar) from one pressure unit to another, in situ if array.

   arguments:
      a (numpy float array, or float): array of pressure values to undergo unit conversion in situ, or a scalar
      from_units (string): the units of the data before conversion
      to_units (string): the required units

   returns:
      a after unit conversion

   note:
      To see supported units, use: `valid_uoms(quantity='pressure')`
   """
   return convert(a, from_units, to_units, quantity = 'pressure', inplace = True)


def convert_volumes(a, from_units, to_units):
   """Converts values in numpy array (or a scalar) from one volume unit to another, in situ if array.

   arguments:
      a (numpy float array, or float): array of volume values to undergo unit conversion in situ, or a scalar
      from_units (string): units of the data before conversion; see note for accepted units
      to_units (string): the required units; see note for accepted units

   returns:
      a after unit conversion

   note:
      To see supported units, use: `valid_uoms(quantity='volume')`

   """
   return convert(a, from_units, to_units, quantity = 'volume', inplace = True)


def convert_flow_rates(a, from_units, to_units):
   """Converts values in numpy array (or a scalar) from one volume flow rate unit to another, in situ if array.

   arguments:
      a (numpy float array, or float): array of volume flow rate values to undergo unit conversion in situ, or a scalar
      from_units (string): units of the data before conversion, eg. 'm3/d'; see notes for acceptable units
      to_units (string): required units of the data after conversion, eg. 'ft3/d'; see notes for acceptable units

   returns:
      a after unit conversion

   note:
      To see supported units, use: `valid_uoms(quantity='volume per time')`
   """
   return convert(a, from_units, to_units, quantity = 'volume per time', inplace = True)


@lru_cache(None)
def get_conversion_factors(uom):
   """Return base unit and conversion factors (A, B, C, D) for a given uom.
   
   The formula "y=(A + Bx)/(C + Dx)" where "y" represents a value in the base unit.

   Returns:
      3-tuple of (base_unit, dimension, factors). Factors is a 4-tuple of conversion factors

   Raises:
      ValueError if either uom is not a valid resqml uom
   """
   if uom not in valid_uoms():
      raise ValueError(f"{uom} is not a valid uom")
   uoms_data = _properties_data()["units"][uom]

   dimension = uoms_data["dimension"]
   try:
      a, b, c, d = uoms_data["A"], uoms_data["B"], uoms_data["C"], uoms_data["D"]
      base_unit = uoms_data["baseUnit"]
   except KeyError:  # Base units do not have factors defined
      a, b, c, d = 0, 1, 1, 0
      base_unit = uom
   return base_unit, dimension, (a, b, c, d)


# Private functions


@lru_cache(None)
def _properties_data():
   """ Return a data structure that represents resqml unit system.

   The dict is loaded directly from a JSON file which is bundled with resqpy.
   The unit system is represented as a dict with the following keys:

   - dimensions
   - quantities
   - units
   - prefixes
   - property_kinds

   Returns:
      dict: resqml unit system
      
   """
   json_path = Path(__file__).parent / 'olio/data/properties.json'
   with open(json_path) as f:
      data = json.load(f)
   return data


def _try_parse_unit(units):
   """Try to match unit against known uoms and aliases, else return None"""

   uom_list = valid_uoms()
   ul = units.casefold()

   if units in uom_list:
      uom = units
   elif ul in CASE_INSENSITIVE_UOMS:
      uom = ul
   elif ul in UOM_ALIAS_MAP:
      uom = UOM_ALIAS_MAP[ul]
   elif ul in uom_list:
      uom = ul  # dangerous! for example, 'D' means D'Arcy and 'd' means day
   else:
      uom = None
   return uom
