BeginPackage["Graphics`OscillatorPlot`"]

Needs["Utilities`FilterOptions`"]

Unprotect[OscillatorPlot,NumericalRegion,NOscillatorCoefficient];
Clear[OscillatorPlot,NumericalRegion,NOscillatorCoefficient];

OscillatorPlot::usage = "OscillatorPlot[f[x],x,{n1,n2}] plots
the energy representation of a given function in the basis of
harmonic oscillator eigenfunctions. The expansion coefficients
are determined by a numerical integration (over an interval
specified by the option NumericalRegion).
The graph shows the expansion coefficients c_n
for n between n1 and n2. The arguments n1 and n2 are optional.
Default values are: n1=0, n2=10."

NOscillatorCoefficient::usage = "NOscillatorCoefficient[expr, x, n]
gives the n-th numerical coefficient in the expansion of the series
expansion of expr (as a function of x) in the basis of eigenfunctions
of the harmonic oscillator."

NumericalRegion::usage = "NumericalRegion is an option for OscillatorPlot
and NOscillatorCoefficient. It specifies the region for
the numerical x-Integration in the calculation of the expansion
coefficients. Default value is NumericalRegion->{-Infinity"

Begin["`Private`"]

OscillatorFunction[n_,x_] :=
   HermiteH[n,x]*Exp[-x^2/2]/Sqrt[2^n*n!]/Pi^(1/4)

Options[OscillatorPlot] = Join[Options[Graphics],
   Options[NIntegrate],{PlotStyle->Thickness[0.01],
   NumericalRegion->{-Infinity,Infinity}}];

SetAttributes[OscillatorPlot,HoldAll];

OscillatorPlot[f_, x_Symbol, opts___Rule] :=
   OscillatorPlot[f, x, {0, 10}, opts]

OscillatorPlot[f_, x_Symbol, {n1_Integer, n2_Integer}, opts___Rule] := 
  Module[{func, coeff, coloredLine, lineTable, up, upper, lower},
    func = Function[{x}, f];
    style = PlotStyle/.{opts}/.Options[OscillatorPlot]; 
    Do[coeff[n] = NOscillatorCoefficient[f, x, n, opts], {n, n1, n2}]; 
    coloredLine[n_] := {style, Hue[Arg[coeff[n]]/(2*Pi)], 
      Line[{{n + 1/2, Abs[coeff[n]]}, 
        {n + 1/2, 0.}}]};
    lineTable = 
     Graphics[Table[coloredLine[n], {n, n1, n2}]]; 
    up = Max[Table[Abs[coeff[n]], {n, n1, n2}]]; 
    upper = up + up/10; lower = -(up/10); 
    Show[lineTable, Evaluate[FilterOptions[Graphics,opts]], Frame -> True,
        PlotRange -> {{n1, n2 + 1}, {lower, upper}}]
  ]
  
Options[NOscillatorCoefficient] = Join[Options[NIntegrate],
   {NumericalRegion->{-Infinity,Infinity}}]

SetAttributes[NOscillatorCoefficient,HoldAll]

NOscillatorCoefficient[f_, x_Symbol, n_, opts___Rule] :=
   Module[{func = Function[{x}, f],
      region = NumericalRegion/.{opts}/.Options[NOscillatorCoefficient]},
      NIntegrate[OscillatorFunction[n, x]*func[x],
         Evaluate[Flatten[{x, region}]],
         Evaluate[FilterOptions[NIntegrate,opts]]
      ]
   ]

End[]

Protect[OscillatorPlot,NumericalRegion,NOscillatorCoefficient]

EndPackage[]