% Compute field propagation matrix
%
% [fDC, fAC, mechTF, tAF] = computeFields(opt, pos, f)
% opt - Optickle model
% pos - optic positions (Nopt x 1, or empty for zeros)
% f - audio frequency vector (Naf x 1)
%
% fDC - DC fields at this position (Nlnk x Nrf)
% fAC - transfer matrix (Narf x Nopt x Naf), where Naff = 2 * Nlnk * Nrf
%   fAC is  arranged such that fAC(n, m, :) is the TF from the
%   position of optic m to DOF n.  The fields are upper and lower
%   audio SBs for each RF component.
% mechTF - mechanical transfer functions (Nopt x Nopt x Naf)
%
% Example:
% f = logspace(0, 3, 300);
% opt = optFP;
% [fDC, fAC, mechTF] = compute(opt, [], f);

% tAF and mMech are reused to make fAC and mechTF
% to avoid memory problems.

function [fDC, tAF, mMech] = computeFields(opt, pos, f)

  % === Get Field Source Info
  [vFrf, vArf, mSrc] = getFieldSource(opt);
  
  % ==== Sizes of Things
  Nopt = opt.Noptic;			% number of optics
  Nlnk = opt.Nlink;			% number of links
  Nprb = opt.Nprobe;			% number of probes
  Nrf  = length(vFrf);			% number of RF components
  Narf = 2 * Nlnk * Nrf;		% number of audio field components
  Ndof = Narf + Nopt;			% number of degrees-of-freedom
  Naf  = length(f);			% number of audio frequencies

  % === Argument Handling
  % set position default
  if isempty(pos)
    pos = zeros(Nopt, 1);
  end
  if length(pos) ~= Nopt
    error('%d positions given, but there are %d optics', length(pos), Nopt);
  end

  % decide if AC calculation is necessary
  isAC = ~isempty(f) && nargout > 1;
  
  % check the memory requirements
  memReq = (20 * Ndof *  Nopt *  Naf) / 1e6;
  if memReq > 200
    qstr = sprintf('This will require about %.0f Mb of memory.', memReq);
    rstr = questdlg([qstr ' Continue?'], 'ComputeFields', 'Yes', 'No', 'No');
    if strcmp(rstr, 'No')
      error('Too much memory.  Exiting.');
    end
  end

  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
  % ==== Frequency Independent Calculations
  % This is where the Optickle model is converted from links
  % and optics to matrices.  The main loop for this section
  % is over all links, since they define the set of fields.
  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

  % initialize matrices
  vDist = zeros(Nlnk, 1);
  mOpt = zeros(Nlnk, Nlnk);
  if Naf > 0
    mMom = zeros(Nopt, Nlnk);
    mGen = zeros(Nlnk, Nlnk, Nopt); % cannot be sparse (more than 2D)
  end
  
  for n = 1:Nlnk
    % prepare some short-cuts
    lnk = opt.link(n);
    src = opt.optic(lnk.snSource);
    snk = opt.optic(lnk.snSink);
    nso = find(snk.out);
    jso = snk.out(nso);

    % ==== Distance Matrix: vDist
    % vDist(n) is the propagation phases-per-frequency
    % from the input to the output of each link in rad/Hz
    vDist(n) = 2 * pi * i  * lnk.len / opt.c;
  
    % ==== Optical Coupling Matrix: mOpt
    % the optical coupling matrix represents the connections
    % between input and output fields provided by each optic,
    % along with the phase from the optic's position
    mOpt(jso, n) = snk.m(nso, lnk.portSink) .* ...
        exp(i * snk.d(nso, lnk.portSink) * opt.k * ...
            (pos(lnk.snSink) + opt.optic(lnk.snSink).pos));

    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% AC calcualtions
    if isAC
      % ==== Momentum Transfer Matrix: mMom
      % the momentum transfer matrix contains the coefficients of
      % momentum transfer from each output field to each optic position
      % this is used to produce radiation pressure effects (see mRad)
      % units are N / W = N / (N * m / s) = s / m
      if snk.tn == opt.types.Mirror | snk.tn == opt.types.BeamSplitter
        mMom(lnk.snSink, n) = sum(abs(mOpt(jso, n)) .*  ...
                                  snk.d(nso, lnk.portSink)) / opt.c; 
      end
      
      % ==== Generation Matrix: mGen
      % the generation matrix represents the phase modulation which
      % generates audio SBs on the input fields (see mMod, below)
      if snk.tn == opt.types.Mirror | snk.tn == opt.types.BeamSplitter
        mGen(jso, n, lnk.snSink) = i * mOpt(jso, n) .* ...
            snk.d(nso, lnk.portSink) * opt.k;
      elseif snk.tn == opt.types.Modulator
        mGen(jso, n, lnk.snSink) = snk.d(nso, lnk.portSink);
      end
      if src.tn == opt.types.Source
        mGen(n, n, lnk.snSource) =  i;
      end      
    end
  end

  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
  % ==== Audio Frequency Independent Calculations
  % Here, the DC fields (fDC) for each RF component, and the
  % matrices which depend on RF frequency, but not audio frequency,
  % are computed (mRad and mMom).  The main loop for this section
  % is over all RF components.
  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

  fDC = zeros(Nlnk, Nrf);			% DC fields
  mRad = sparse(Nopt, Narf);			% radiation force coupling
  mMod = sparse(Narf, Nopt);			% position ASB coupling

  jLnk = 1:Nlnk;
  eyeNlnk = speye(Nlnk);
  mPhi = zeros(Nlnk, Nlnk);			% propagation phase matrix
  for n = 1:Nrf
    % indices for the upper (plus) and lower (minus) audio SBs
    np = 2 * Nlnk * (n - 1) + jLnk;
    nm = np + Nlnk;
    
    % ==== DC Fields: fDC  (Nlnk x Nrf)
    % fDC is a vector of fields from DC sources
    % these fields are the seed fields for all audio SBs
    % and radiation pressure couplings (see mMod and mRad, below)
    mPhi = diag(exp(vDist * vFrf(n)));
    fDCn = (eyeNlnk - (mPhi * mOpt)) \ (mPhi * mSrc(:, n));
    mDCn = diag(fDCn);				% matrix of output fields
    fDC(:, n) = fDCn;				% record for later
    
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% AC calcualtions
    if isAC
      % ==== Radiation Forces: mRad (Nopt x Narf)
      % mRad is the transfer matrix from audio SB amplitude to force
      % the units are N / sqrt(W).
      mRad(:, nm) = mMom * conj(mDCn);
      mRad(:, np) = conj(mMom * mDCn);		% conj operator for upper SB

      % ==== Modulation: mMod (Narf x Nopt)
      % mMod is the transfer matrix from optic position to audio SB amplitude
      for m = 1:Nopt
        mMod(nm, m) = mGen(:, :, m) * fDCn / 2;
      end
      mMod(np, :) = conj(mMod(nm, :));		% conj operator for upper SB
    end
  end

  % if AC is not needed, just end here
  if ~isAC
    tAF = zeros(Narf, Nopt, 0);
    mMech = zeros(Nopt, Nopt, 0);
    return
  end
  
  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
  % ==== Audio Frequency Loop
  % This is the section in which the tranfer functions
  % from optic position to fields (and optic positions) are
  % computed (tAF).  tAF is Ndof x Nopt x Naf and arranged
  % such that tAF(n, m, :) is the TF from the position of optic m
  % to DOF n.  The DOFs are upper and lower audio SBs for each RF
  % component, followed by the optic's positions.  The fields are
  % only computed at the output of each link.
  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

  tAF = zeros(Ndof, Nopt, Naf);
  
  mDof = sparse(Ndof, Ndof);
  jFld = 1:Narf;
  jOpt = (1:Nopt) + Narf;
  
  % excitation matrix: 1 for each optic
  eyeNdof = speye(Ndof);
  mExc = eyeNdof(:, jOpt);
  
  % mechanical transfer matrix
  mMech = zeros(Nopt, Nopt, Naf);
  for nOpt = 1:Nopt
    if ~isempty(opt.optic(nOpt).mechTF)
      rsp = freqresp(opt.optic(nOpt).mechTF, 2 * pi * f);
      if ndims(rsp) == 1
        mMech(nOpt, nOpt, :) = rsp;
      else
        mMech(nOpt, nOpt, :) = rsp(1, 1, :);
      end
    end
  end

  % since this can take a while, let's time it
  tic;
  hWaitBar = [];
  tLast = 0;
  
  % audio frequency loop
  for nAF = 1:Naf
    fAudio = f(nAF);

    % ==== Degree-of-Freedom Coupling: mDof (Ndof x Ndof)

    % field to optic position transfer (only at link outputs = even fields)
    mDof(jOpt, jFld) = mMech(:, :, nAF) * mRad(:, :);

    % field to field transfer
    for n = 1:Nrf
      np = 2 * Nlnk * (n - 1) + jLnk;
      nm = np + Nlnk;

      % upper audio frequency
      mPhi = diag(exp(-vDist * (vFrf(n) + fAudio)));
      mDof(np, np) = mPhi * conj(mOpt);		% conj operator for upper SB
      mDof(np, jOpt) = mPhi * mMod(np, :);      % optic position to field

      % lower audio frequency
      mPhi = diag(exp(vDist * (vFrf(n) - fAudio)));
      mDof(nm, nm) = mPhi * mOpt;
      mDof(nm, jOpt) = mPhi * mMod(nm, :);      % optic position to field
    end
  
    % solve (yes, this is where the magic happens)
    tAF(:, :, nAF) = (eyeNdof - mDof) \ mExc;

    % ==== Timing and User Interaction
    % NO MODELING HERE (just let the user know how long this will take)
    tNow = toc;
    frac = nAF / Naf;
    tRem = tNow * (1 / frac - 1);
    if tNow > 2 & tRem > 2 & tNow - tLast > 0.5
      % wait bar string
      str = sprintf('Computing fields:  %.1f s used, %.1f s left', ...
                     tNow, tRem);

      % check and update waitbar
      if isempty(hWaitBar)
        % create wait bar
        try
          str = [str ' (close this window to stop)'];
          hWaitBar = waitbar(frac, str);
          tLast = tNow;
        catch
          % can't make wait bar... use text
          if tNow - tLast > 5
            disp(str)
            tLast = tNow;
          end
        end
      else
        try
          str = [str ' (close this window to stop)'];
          findobj(hWaitBar);			% error if wait bar is closed
          waitbar(frac, hWaitBar, str);		% update wait string
          tLast = tNow;
        catch
          error('Wait bar closed by user.  Exiting.')
        end
      end
    end
  end
    
  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
  % ==== Done
  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

  % close wait bar
  if ~isempty(hWaitBar)
    waitbar(1.0, hWaitBar, 'Done computing fields.  Returning...')
    close(hWaitBar)
  end

  % mechanical part of tAF (should be mechTF)
  mMech = tAF(jOpt, :, :);

  % optical part of tAF (should be fAC)
  tAF = tAF(jFld, :, :);

  % make sure that the wait bar is closed
  drawnow

  return
  
  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
  % ==== Shot Noise Loop
  % Here the transfer function from each source and sink output
  % is copied into TFs to each port.  This can later be used for
  % shot noise computation.
  %
  % This doesn't work because tAF does not include the field
  % inputs.  This computation should be moved into the main audio
  % loop, and it will probably cause memory problems. =(
  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

  % count sources and sinks
  jShot = [];
  Nshot = 0;
  for n = 1:Nopt
    % prepare some short-cuts
    src = opt.optic(n);

    if src.tn == opt.types.Source | src.tn == opt.types.Sink
      Nshot = Nshot + 1;
      jShot(Nshot) = n;
    end
  end
  
  % create the shot-noise transfer matrix
  mShot = zeros(Ndof, 2 * Nrf, Naf, Nshot);
  
  jRF = 1:Nlnk:Narf;
  for n = 1:Nshot
    % prepare some short-cuts
    src = opt.optic(jShot(n));
    jno = src.out(1) + jRF;
  
    % fill in this part of the matrix
    mShot(:, :, :, n) = tAF(:, jno, :);
  end
  
