(************************************************************************)
(* Package for natural unit conversions.                                *)
(************************************************************************)

BeginPackage["NaturalUnits`"];

Print["NaturalUnits 1.0"];

(************ Interface **************)

(* Functions and setting variables *)
NaturalConvert::usage="Convert quantities with units into other units, while setting c=hbar=kB=1. First parameter: quantity to be converted multiplied by unit; second parameter: target unit (default: toUnitDefault=nounit); third parameter: keep units in the final result or not (True, False or Defer; default: keepUnitDefault=True).";
toUnitDefault::usage="Default value of the second parameter of NaturalConvert. Starts as nounit, corresponding to powers of eV.";
keepUnitDefault::usage="Default value of the third parameter of NaturalConvert. Starts as True.";
nounit::usage="Defaul unit: corresponds to some power of eV.";

NaturalUnitsCheck::usage="Check if the symbol given as a parameter is defined as a unit symbol.";

NaturalUnitsSetup::usage="Set up substitution rules and other expressions for unit manioulations. Called when package is loaded; should be called again if basicUnits, derivedUnits, prefixableUnits, or prefixTable is changed.";

(* Unit lists *)
basicUnits::usage="Replacement rules for basic SI units.";
derivedUnits::usage="Replacement rules for units derived from other units.";
prefixableUnits::usage="Units that take metric prefixes.";
prefixTable::usage="List of metric prefixes.";
toNatUnits::usage="Replacement rule collecting all other repalcement rules (basicUnits, derivedUnits, prefixedUnits).";
unitList::usage="List of all unit symbols.";

(* Base units and constants *)
m::usage="meter";
s::usage="second";
kg::usage="kilogram";
K::usage="Kelvin";
min::usage="minute";
year::usage="year";
Hz::usage="Hertz";
mP::usage="Planck mass";
MP::usage="reduced Planck mass";
tP::usage="Planck time";
TP::usage="Planck temperature";
lP::usage="Planck length";
EP::usage="Planck energy";
pc::usage="parsec";
ly::usage="lightyear";
AU::usage="astronomical unit";
g::usage="gram";
Msolar::usage="solar mass";
MEarth::usage="Earth mass";
MMoon::usage="Moon mass";
J::usage="Joule";
h::usage="hour";
month::usage="month";
y::usage="year";

c::usage="speed of light";
hbar::usage="reduced Planck constant";
kB::usage="Boltzmann constant";
G::usage="Newton's constant";
eV::usage="electronvolt";
natUnit::usage="natural unit: eV";

(************ Conversion data **************)

(* Constants needed to fix the natural units, in terms of basic units. *)
c = 299792458 m/s;
hbar = 6.62607015*10^-34 m^2*kg/s/(2*Pi);
kB = 1.380649*10^-23 kg*m^2/s^2/K;
G = 6.674*10^-11 m^3/(kg*s^2);
natUnit=eV;

(* Basic units solved in terms of the used natural unit. *)
basicUnits = Association[Quiet[Solve[{c == 1, hbar == 1, kB == 1, natUnit == 1.602176634*10^-19 kg*m^2/s^2}, {m, s, kg, K}]]];

(* Derived units in terms of the basic units. This is a list of associations. Substitution rules in later associations can use units defined in former associations. Add new units to the end. *)
derivedUnits = {
<|min->60 s,
year -> 31556926 s,
Hz -> 1/s,
mP -> Sqrt[hbar*c/G],
MP -> Sqrt[hbar*c/(8*Pi*G)],
tP -> Sqrt[hbar*G/c^5],
TP -> Sqrt[hbar*c^5/(G*kB^2)],
lP -> Sqrt[hbar*G/c^3],
EP->Sqrt[hbar*c^5/G],
pc -> 3.08567758*10^16 m,
ly->9.4605284*10^15 m,
AU -> 149597871*10^3 m,
g->10^-3 kg,
Msolar->1.989*10^30 kg,
MEarth->5.972*10^24 kg,
MMoon->7.348*10^22 kg,
J->kg*m^2/s^2|>,
<|h->60 min,
month->year/12,
y->year|>
};

(* Units that take matric prefixes (milli-, centi-, kilo-, ... *)
prefixableUnits={s,y,Hz,m,pc,ly,eV,J,g};

(* Metric prefixes and the corresponding multipliers. *)
prefixTable=<|"mu"->10^-6,"m"->10^-3,"c"->10^-2,"k"->10^3,"M"->10^6,"G"->10^9|>;

(* To be set later. *)
toNatUnits=<||>;

(************ Messages (used in functions below) **************)
NatAssociateTo::unitdpl="Some units duplicated in `` and ``. Duplicated units and tallies: ``";
NaturalConvert::npowers="Error: Parameters contain multiple incompatible powers of natural units.";
NaturalConvert::unitcomp="Units of `` and `` not compatible! (Or unit not recognized)";
NaturalConvert::corrpow="Trying to correct the power of target units.";

(************ Setup functions (public implementation: add new public variables) **************)

(* Custom function for appending associations which warns for repeated keys. *)
SetAttributes[NatAssociateTo, HoldFirst];
NatAssociateTo[asTo_, asFrom_]:=Module[{duplicates=Select[Tally[Join[Keys[asTo],Keys[asFrom]]], (#[[2]]>1)&]},
If[Length[duplicates]>0, Message[NatAssociateTo::unitdpl, asTo, asFrom, duplicates]];
AssociateTo[asTo,asFrom];
];



(* Parse unit lists (basicUnits, derivedUnits, prefixableUnits, prefixTable) into usable form (toNatUnits). Called once when the package is loaded. Call again, if lists changed on the fly. *)
NaturalUnitsSetup[]:=Module[{newUnits, newSymbol},

(* Join basic and derived units into one association in terms of natUnits. *)
toNatUnits=basicUnits;
Map[Function[{currentTable}, NatAssociateTo[toNatUnits,currentTable/.toNatUnits]], derivedUnits];

(* Do some manual adjustments before adding the prefixes. *)
NatAssociateTo[toNatUnits,<|natUnit->natUnit|>];
toNatUnits=KeyDrop[toNatUnits,kg]; (* Already have grams; prefixed mass units to be constructed from g instead of kg. *)

(* Add prefixes to those units that need them. *)
Map[Function[{unitNow},
newUnits=Map[Function[{prefixNow}, newSymbol=Symbol[prefixNow<>ToString[unitNow]];(Evaluate[newSymbol]::usage=ToString[StringForm["unit formed from `` by prefix ``",unitNow::usage,prefixNow]]);newSymbol->prefixTable[prefixNow]*toNatUnits[unitNow]], Keys[prefixTable]];
NatAssociateTo[toNatUnits,newUnits]], prefixableUnits];

(* Finalize variables. *)
toNatUnits=AssociationMap[N[Simplify[#,natUnit>0]]&, toNatUnits];
unitAssumptions=And@@Map[(#>0)&,Keys[toNatUnits]];
unitList=Keys[toNatUnits];

];
NaturalUnitsSetup[]


(************ Other functions (private implementation) **************)

Begin["`Private`"];

(* Check whether a given unit exists. *)
NaturalUnitsCheck[checkedUnit_]:=Module[{},
If[KeyExistsQ[toNatUnits, checkedUnit],
Print[StringForm["Unit `` exists as ``=``", checkedUnit, checkedUnit, toNatUnits[checkedUnit]]];
If[KeyExistsQ[(Join@@derivedUnits), checkedUnit], Print[StringForm["Original definition: ``=``", checkedUnit, (Join@@derivedUnits)[checkedUnit]]]],
Print[StringForm["Unit `` does not exist.", checkedUnit]];
];
];



(* Convert one unit to another. *)
NaturalConvert[qty_, toUnitIn_:toUnitDefault, keepUnit_:keepUnitDefault]:=Module[{toUnit=toUnitIn, toUnitDef=toUnitIn, qtyNat, toUnitNat, toUnitExp, qtyFinal},

(* First, convert quantity and target unit to natural units. *)
qtyNat=qty/.toNatUnits;
toUnitNat=toUnit/.toNatUnits;

(* If natural units wanted, go for them. *)
If[toUnit===nounit, Return[If[keepUnit===False,(qtyNat/.natUnit->1),FullSimplify[qtyNat,natUnit>0]]]];

(* If quantity or target units are not monomial w.r.t. the natural units, reject this case. *)
If[Quiet[OnePowerQ[qtyNat,natUnit]]==False || Quiet[OnePowerQ[toUnitNat,natUnit]]==False,
Message[NaturalConvert::npowers];
Return[qty];
];

(* Check if the units match. If not, try to correct this. *)
If[Exponent[qtyNat,natUnit]!=Exponent[toUnitNat,natUnit],
Message[NaturalConvert::unitcomp, qty, toUnit];

If[Exponent[qtyNat,natUnit]!=0 && Exponent[toUnitNat,natUnit]!=0,
Message[NaturalConvert::corrpow];
toUnitExp=(Exponent[qtyNat,natUnit]/Exponent[toUnitNat,natUnit]);
toUnitDef=Defer["toUnit"^"toUnitExp"]/.{"toUnit"->toUnit,"toUnitExp"->toUnitExp};
toUnit=Simplify[toUnit^toUnitExp,unitAssumptions];
toUnitNat=toUnitNat^toUnitExp,
Return[qty];
]];

(* Calculate the final number. *)
qtyFinal=qtyNat/toUnitNat/.natUnit->1; (* Set natUnit->1 to get rid of any unexpected errors. *)

(* Add the unit, if wanted. *)
If[keepUnit===False,qtyFinal,
If[keepUnit===Defer,(Defer["qty"*"toUnit"]/.{"qty"->qtyFinal,"toUnit"->toUnitDef}),
qtyFinal*toUnit]]
];
toUnitDefault=nounit;
keepUnitDefault=True;



(* Helper function: Check if expression contains only one power of given variable (to the best of our ability). *)
OnePowerQ[expr_,x_]:=Module[{testExpr=Simplify[x*D[expr,x]/expr,x>0]},
If[testExpr===(testExpr/.x->0),True,False]
];

End[];

EndPackage[];
