function [ cluster, cluster_center, cluster_population, cluster_energy, ...
  it_num ] = kmeans_02 ( dim_num, point_num, cluster_num, it_max, point, ...
  cluster_center )

%*****************************************************************************80
%
%% KMEANS_02 applies the K-Means algorithm.
%
%  Discussion:
%
%    The routine attempts to divide POINT_NUM points in
%    DIM_NUM-dimensional space into CLUSTER_NUM clusters so that the within
%    cluster sum of squares is minimized.
%
%  Licensing:
%
%    This code is distributed under the GNU LGPL license.
%
%  Modified:
%
%    04 October 2009
%
%  Author:
%
%    Original FORTRAN77 by John Hartigan, Manchek Wong.
%    FORTRAN90 version by John Burkardt.
%
%  Reference:
%
%    John Hartigan, Manchek Wong,
%    Algorithm AS 136:
%    A K-Means Clustering Algorithm,
%    Applied Statistics,
%    Volume 28, Number 1, 1979, pages 100-108.
%
%  Parameters:
%
%    Input, integer DIM_NUM, the number of spatial dimensions.
%
%    Input, integer POINT_NUM, the number of points.
%
%    Input, integer CLUSTER_NUM, the number of clusters.
%
%    Input, integer IT_MAX, the maximum number of iterations.
%
%    Input, real POINT(DIM_NUM,POINT_NUM), the coordinates
%    of the points.
%
%    Input, real CLUSTER_CENTER(DIM_NUM,CLUSTER_NUM),
%    the cluster centers.
%
%    Output, integer CLUSTER(POINT_NUM), the cluster each
%    point belongs to.
%
%    Output, real CLUSTER_CENTER(DIM_NUM,CLUSTER_NUM),
%    the cluster centers.
%
%    Output, integer CLUSTER_POPULATION(CLUSTER_NUM), the number
%    of points in each cluster.
%
%    Output, real CLUSTER_ENERGY(CLUSTER_NUM), the
%    within-cluster sum of squares.
%
%    Output, integer IT_NUM, the number of iterations taken.
%
  it_num = 0;

  if ( cluster_num < 1 )
    fprintf ( 1, '\n' );
    fprintf ( 1, 'KMEANS_02 - Fatal error!\n' );
    fprintf ( 1, '  CLUSTER_NUM < 1.\n' );
    error ( 'KMEANS_02 - Fatal error!' )
  end

  if ( point_num <= cluster_num )
    fprintf ( 1, '\n' );
    fprintf ( 1, 'KMEANS_02 - Fatal error!\n' );
    fprintf ( 1, '  POINT_NUM <= CLUSTER_NUM.\n' );
    error ( 'KMEANS_02 - Fatal error!' )
  end

  if ( dim_num < 1 )
    fprintf ( 1, '\n' );
    fprintf ( 1, 'KMEANS_02 - Fatal error!\n' );
    fprintf ( 1, '  DIM_NUM < 1.\n' );
    error ( 'KMEANS_02 - Fatal error!' )
  end

  if ( point_num < 1 )
    fprintf ( 1, '\n' );
    fprintf ( 1, 'KMEANS_02 - Fatal error!\n' );
    fprintf ( 1, '  POINT_NUM < 1.\n' );
    error ( 'KMEANS_02 - Fatal error!' )
  end
%
%  For each point I, find its two closest centers, CLUSTER(I) and CLUSTER2(I).
%  Assign it to CLUSTER(I).
%
  for i = 1 : point_num

    cluster(i) = 1;
    cluster2(i) = 2;

    for il = 1 : 2
      dt(il) = sum ( ...
        ( point(1:dim_num,i) - cluster_center(1:dim_num,il) ).^2 );
    end

    if ( dt(2) < dt(1) )
      cluster(i) = 2;
      cluster2(i) = 1;
      temp = dt(1);
      dt(1) = dt(2);
      dt(2) = temp;
    end

    for l = 3 : cluster_num

      db = sum ( ( point(1:dim_num,i) - cluster_center(1:dim_num,l) ).^2 );

      if ( db < dt(1) )
        dt(2) = dt(1);
        cluster2(i) = cluster(i);
        dt(1) = db;
        cluster(i) = l;
      elseif ( db < dt(2) )
        dt(2) = db;
        cluster2(i) = l;
      end

    end

  end
%
%  Update cluster centers to be the average of points contained within them.
%
  cluster_population(1:cluster_num) = 0;
  cluster_center(1:dim_num,1:cluster_num) = 0.0;

  for i = 1 : point_num
    l = cluster(i);
    cluster_population(l) = cluster_population(l) + 1;
    cluster_center(1:dim_num,l) = cluster_center(1:dim_num,l) ...
      + point(1:dim_num,i);
  end
%
%  Check to see if there is any empty cluster.
%
  for l = 1 : cluster_num

    if ( cluster_population(l) == 0 )
      j = 1 + floor ( point_num * rand ( ) );
      cluster_center(1:dim_num) = point(1:dim_num,j)
    else
      cluster_center(1:dim_num,l) = cluster_center(1:dim_num,l) ...
        / cluster_population(l);
    end
%
%  Initialize AN1, AN2, ITRAN and NCP
%  AN1(L) = CLUSTER_POPULATION(L) / (CLUSTER_POPULATION(L) - 1)
%  AN2(L) = CLUSTER_POPULATION(L) / (CLUSTER_POPULATION(L) + 1)
%  ITRAN(L) = 1 if cluster L is updated in the quick-transfer stage,
%           = 0 otherwise
%  In the optimal-transfer stage, NCP(L) stores the step at which
%  cluster L is last updated.
%  In the quick-transfer stage, NCP(L) stores the step at which
%  cluster L is last updated plus POINT_NUM.
%
    an2(l) = cluster_population(l) / ( cluster_population(l) + 1 );

    if ( 1 < cluster_population(l) )
      an1(l) = cluster_population(l) / ( cluster_population(l) - 1 );
    else
      an1(l) = Inf;
    end

    itran(l) = 1;
    ncp(l) = -1;

  end

  indx = 0;
  ifault = 2;
  it_num = 0;

  d = zeros(point_num);
  live = zeros(cluster_num);

  while ( it_num < it_max )

    it_num = it_num + 1;
%
%  In this stage, there is only one pass through the data.   Each
%  point is re-allocated, if necessary, to the cluster that will
%  induce the maximum reduction in within-cluster sum of squares.
%
    [ cluster_center, cluster, cluster2, cluster_population, an1, an2, ...
      ncp, d, itran, live, indx ] = kmeans_02_optra ( dim_num, point_num, ...
      cluster_num, point, cluster_center, cluster, cluster2, cluster_population, ...
      an1, an2, ncp, d, itran, live, indx );
%
%  Stop if no transfer took place in the last POINT_NUM optimal transfer steps.
%
    if ( indx == point_num )
      ifault = 0;
      break
    end
%
%  Each point is tested in turn to see if it should be re-allocated
%  to the cluster to which it is most likely to be transferred,
%  CLUSTER2(I), from its present cluster, CLUSTER(I).   Loop through the
%  data until no further change is to take place.
%
    [ cluster_center, cluster, cluster2, cluster_population, an1, ...
      an2, ncp, d, itran, indx ] = kmeans_02_qtran ( dim_num, point_num, ...
      cluster_num, point, cluster_center, cluster, cluster2, cluster_population, ...
      an1, an2, ncp, d, itran, indx );
%
%  If there are only two clusters, there is no need to re-enter the
%  optimal transfer stage.
%
    if ( cluster_num == 2 )
      ifault = 0;
      break
    end
%
%  NCP has to be set to 0 before entering OPTRA.
%
    ncp(1:cluster_num) = 0;

  end

  if ( ifault == 2 )
    fprintf ( 1, '\n' );
    fprintf ( 1, 'KMEANS_02 - Warning!\n' );
    fprintf ( 1, '  Maximum number of iterations reached\n' );
    fprintf ( 1, '  without convergence.\n' );
  end
%
%  Compute the within-cluster sum of squares for each cluster.
%
  cluster_center(1:dim_num,1:cluster_num) = 0.0;

  for i = 1 : point_num
    cluster_center(1:dim_num,cluster(i)) = ...
      cluster_center(1:dim_num,cluster(i)) + point(1:dim_num,i);
  end

  for j = 1 : dim_num
    cluster_center(j,1:cluster_num) = cluster_center(j,1:cluster_num) ...
      ./ cluster_population(1:cluster_num);
  end
%
%  Compute the cluster energies.
%
  cluster_energy(1:cluster_num) = 0.0;

  for i = 1 : point_num

    j = cluster(i);

    cluster_energy(j) = cluster_energy(j) + sum ( ...
      ( point(1:dim_num,i) - cluster_center(1:dim_num,j) ).^2 );

  end

  return
end
function [ cluster_center, cluster, cluster2, cluster_population, an1, an2, ...
  ncp, d, itran, live, indx ] = kmeans_02_optra ( dim_num, point_num, ...
  cluster_num, point, cluster_center, cluster, cluster2, cluster_population, ...
  an1, an2, ncp, d, itran, live, indx )

%*****************************************************************************80
%
%% KMEANS_02_OPTRA carries out the optimal transfer stage.
%
%  Discussion:
%
%    Each point is re-allocated, if necessary, to the cluster that
%    will induce a maximum reduction in the within-cluster sum of
%    squares.
%
%  Licensing:
%
%    This code is distributed under the GNU LGPL license.
%
%  Modified:
%
%    04 October 2009
%
%  Author:
%
%    Original FORTRAN77 by John Hartigan, Manchek Wong.
%    FORTRAN90 version by John Burkardt.
%
%  Reference:
%
%    John Hartigan, Manchek Wong,
%    Algorithm AS 136:
%    A K-Means Clustering Algorithm,
%    Applied Statistics,
%    Volume 28, Number 1, 1979, pages 100-108.
%
%  Parameters:
%
%    Input, integer DIM_NUM, the number of spatial dimensions.
%
%    Input, integer POINT_NUM, the number of points.
%
%    Input, integer CLUSTER_NUM, the number of clusters.
%
%    Input, real POINT(DIM_NUM,POINT_NUM), the coordinates of
%    the points.
%
%    Input, real CLUSTER_CENTER(DIM_NUM,CLUSTER_NUM),
%    the cluster centers.
%
%    Input, integer CLUSTER(POINT_NUM), the cluster
%    each point belongs to.
%
%    Input, integer CLUSTER2(POINT_NUM), the cluster
%    to which each point is most likely to be transferred to.
%
%    Input, integer CLUSTER_POPULATION(CLUSTER_NUM),
%    the number of points in each cluster.
%
%    Input, real AN1(CLUSTER_NUM),
%    CLUSTER_POPULATION(L) / (CLUSTER_POPULATION(L) - 1)
%
%    Input, real AN2(CLUSTER_NUM),
%    CLUSTER_POPULATION(L) / (CLUSTER_POPULATION(L) + 1)
%
%    Input, integer NCP(CLUSTER_NUM), ?
%
%    Input, real D(POINT_NUM), ?
%
%    Input, integer ITRAN(CLUSTER_NUM),
%    1 if cluster L is updated in the quick-transfer stage,
%    0 otherwise.  Reset to 0 on output.
%
%    Input, integer LIVE(CLUSTER_NUM), ?
%
%    Input, integer INDX, ?
%
%    Output, real CLUSTER_CENTER(DIM_NUM,CLUSTER_NUM),
%    the cluster centers.
%
%    Output, integer CLUSTER(POINT_NUM), the cluster
%    each point belongs to.
%
%    Output, integer CLUSTER2(POINT_NUM), the cluster
%    to which each point is most likely to be transferred to.
%
%    Output, integer CLUSTER_POPULATION(CLUSTER_NUM),
%    the number of points in each cluster.
%
%    Output, real AN1(CLUSTER_NUM),
%    CLUSTER_POPULATION(L) / (CLUSTER_POPULATION(L) - 1)
%
%    Output, real AN2(CLUSTER_NUM),
%    CLUSTER_POPULATION(L) / (CLUSTER_POPULATION(L) + 1)
%
%    Output, integer NCP(CLUSTER_NUM), ?
%
%    Output, real D(POINT_NUM), ?
%
%    Output, integer ITRAN(CLUSTER_NUM),
%    1 if cluster L is updated in the quick-transfer stage,
%    0 otherwise.  Reset to 0 on output.
%
%    Output, integer LIVE(CLUSTER_NUM), ?
%
%    Output, integer INDX, ?
%

%
%  If cluster L is updated in the last quick-transfer stage, it
%  belongs to the live set throughout this stage.   Otherwise, at
%  each step, it is not in the live set if it has not been updated
%  in the last POINT_NUM optimal transfer steps.
%
  for l = 1 : cluster_num
    if ( itran(l) == 1 )
      live(l) = point_num + 1;
    end
  end

  for i = 1 : point_num

    indx = indx + 1;
    l1 = cluster(i);
    l2 = cluster2(i);
    ll = l2;
%
%  If point I is the only member of cluster L1, no transfer.
%
    if ( 1 < cluster_population(l1) )
%
%  If L1 has been updated in this stage, re-compute D(I).
%
      if ( ncp(l1) ~= 0 )
        d(i) = an1(l1) * sum ( ...
          ( point(1:dim_num,i) - cluster_center(1:dim_num,l1) ).^2 );
      end
%
%  Find the cluster with minimum R2.
%
      r2 = an2(l2) * sum ( ...
        ( point(1:dim_num,i) - cluster_center(1:dim_num,l2) ).^2 );

      for l = 1 : cluster_num
%
%  If LIVE(L1) <= I, then L1 is not in the live set.   If this is
%  true, we only need to consider clusters that are in the live set
%  for possible transfer of point I.
%
%  Otherwise, we need to consider all possible clusters.
%
        if ( ( i < live(l1) | i < live(l) ) & ...
             l ~= l1 & ...
             l ~= ll )

          rr = r2 / an2(l);

          dc = sum ( ( point(1:dim_num,i) - cluster_center(1:dim_num,l) ).^2 );

          if ( dc < rr )
            r2 = dc * an2(l);
            l2 = l;
          end

        end

      end
%
%  If no transfer is necessary, L2 is the new CLUSTER2(I).
%
      if ( d(i) <= r2 )

        cluster2(i) = l2;

      else
%
%  Update cluster centers, LIVE, NCP, AN1 and AN2 for clusters L1 and
%  L2, and update CLUSTER(I) and CLUSTER2(I).
%
        indx = 0;
        live(l1) = point_num + i;
        live(l2) = point_num + i;
        ncp(l1) = i;
        ncp(l2) = i;
        al1 = cluster_population(l1);
        alw = al1 - 1.0;
        al2 = cluster_population(l2);
        alt = al2 + 1.0;

        cluster_center(1:dim_num,l1) = ( cluster_center(1:dim_num,l1) * al1 ...
          - point(1:dim_num,i) ) / alw;

        cluster_center(1:dim_num,l2) = ( cluster_center(1:dim_num,l2) * al2 ...
          + point(1:dim_num,i) ) / alt;

        cluster_population(l1) = cluster_population(l1) - 1;
        cluster_population(l2) = cluster_population(l2) + 1;
        an2(l1) = alw / al1;

        if ( 1.0 < alw )
          an1(l1) = alw / ( alw - 1.0 );
        else
          an1(l1) = Inf;
        end

        an1(l2) = alt / al2;
        an2(l2) = alt / ( alt + 1.0 );
        cluster(i) = l2;
        cluster2(i) = l1;

      end

    end

    if ( indx == point_num )
      return
    end

  end
%
%  ITRAN(L) = 0 before entering QTRAN.
%
  itran(1:cluster_num) = 0;
%
%  LIVE(L) has to be decreased by POINT_NUM before re-entering OPTRA.
%
  live(1:cluster_num) = live(1:cluster_num) - point_num;

  return
end
function [ cluster_center, cluster, cluster2, cluster_population, an1, ...
  an2, ncp, d, itran, indx ] = kmeans_02_qtran ( dim_num, point_num, ...
  cluster_num, point, cluster_center, cluster, cluster2, cluster_population, ...
  an1, an2, ncp, d, itran, indx )

%*****************************************************************************80
%
%% KMEANS_02_QTRAN carries out the quick transfer stage.
%
%  Discussion:
%
%    For each point I, CLUSTER(I) and CLUSTER2(I) are switched, if necessary,
%    to reduce within-cluster sum of squares.  The cluster centers are
%    updated after each step.
%
%  Licensing:
%
%    This code is distributed under the GNU LGPL license.
%
%  Modified:
%
%    04 October 2009
%
%  Author:
%
%    Original FORTRAN77 by John Hartigan, Manchek Wong.
%    FORTRAN90 version by John Burkardt.
%
%  Reference:
%
%    John Hartigan, Manchek Wong,
%    Algorithm AS 136:
%    A K-Means Clustering Algorithm,
%    Applied Statistics,
%    Volume 28, Number 1, 1979, pages 100-108.
%
%  Parameters:
%
%    Input, integer DIM_NUM, the number of spatial dimensions.
%
%    Input, integer POINT_NUM, the number of points.
%
%    Input, integer CLUSTER_NUM, the number of clusters.
%
%    Input, real POINT(DIM_NUM,POINT_NUM), the coordinates
%    of the points.
%
%    Input, real CLUSTER_CENTER(DIM_NUM,CLUSTER_NUM),
%    the cluster centers.
%
%    Input, integer CLUSTER(POINT_NUM), the cluster
%    each point belongs to.
%
%    Input, integer CLUSTER2(POINT_NUM), the cluster to which
%    each point is most likely to be transferred to.
%
%    Input, integer CLUSTER_POPULATION(CLUSTER_NUM),
%    the number of points in each cluster.
%
%    Input, real AN1(CLUSTER_NUM),
%    CLUSTER_POPULATION(L) / (CLUSTER_POPULATION(L) - 1)
%
%    Input, real AN2(CLUSTER_NUM),
%    CLUSTER_POPULATION(L) / (CLUSTER_POPULATION(L) + 1)
%
%    Input, integer NCP(CLUSTER_NUM), ?
%
%    Input, real D(POINT_NUM), ?
%
%    Input, integer ITRAN(CLUSTER_NUM),
%    1 if cluster L is updated in the quick-transfer stage,
%    0 otherwise.
%
%    Input, integer INDX, is set to 0 if any updating occurs.
%
%    Output, real CLUSTER_CENTER(DIM_NUM,CLUSTER_NUM),
%    the cluster centers.
%
%    Output, integer CLUSTER(POINT_NUM), the cluster
%    each point belongs to.
%
%    Output, integer CLUSTER2(POINT_NUM), the cluster to which
%    each point is most likely to be transferred to.
%
%    Output, integer CLUSTER_POPULATION(CLUSTER_NUM),
%    the number of points in each cluster.
%
%    Output, real AN1(CLUSTER_NUM),
%    CLUSTER_POPULATION(L) / (CLUSTER_POPULATION(L) - 1)
%
%    Output, real AN2(CLUSTER_NUM),
%    CLUSTER_POPULATION(L) / (CLUSTER_POPULATION(L) + 1)
%
%    Output, integer NCP(CLUSTER_NUM), ?
%
%    Output, real D(POINT_NUM), ?
%
%    Output, integer ITRAN(CLUSTER_NUM),
%    1 if cluster L is updated in the quick-transfer stage,
%    0 otherwise.
%
%    Output, integer INDX, is set to 0 if any updating occurs.
%

%
%  In the optimal transfer stage, NCP(L) indicates the step at which
%  cluster L is last updated.   In the quick transfer stage, NCP(L)
%  is equal to the step at which cluster L is last updated plus POINT_NUM.
%
  count = 0;
  step = 0;

  while ( 1 )

    for i = 1 : point_num

      count = count + 1;
      step = step + 1;
      l1 = cluster(i);
      l2 = cluster2(i);
%
%  If point I is the only member of cluster L1, no transfer.
%
      if ( 1 < cluster_population(l1) )
%
%  If NCP(L1) < STEP, no need to re-compute distance from point I to
%  cluster L1.   Note that if cluster L1 is last updated exactly POINT_NUM
%  steps ago, we still need to compute the distance from point I to
%  cluster L1.
%
        if ( step <= ncp(l1) )
          d(i) = an1(l1) * sum ( ...
            ( point(1:dim_num,i) - cluster_center(1:dim_num,l1) ).^2 );
        end
%
%  If STEP >= both NCP(L1) and NCP(L2) there will be no transfer of
%  point I at this step.
%
        if ( step < ncp(l1) | step < ncp(l2) )

          r2 = d(i) / an2(l2);

          dd = sum ( ( point(1:dim_num,i) - cluster_center(1:dim_num,l2) ).^2 );
%
%  Update cluster centers, NCP, CLUSTER_POPULATION, ITRAN, AN1 and AN2
%  for clusters L1 and L2.   Also update CLUSTER(I) and CLUSTER2(I).
%
%  Note that if any updating occurs in this stage, INDX is set back to 0.
%
          if ( dd < r2 )

            count = 0;
            indx = 0;
            itran(l1) = 1;
            itran(l2) = 1;
            ncp(l1) = step + point_num;
            ncp(l2) = step + point_num;
            al1 = cluster_population(l1);
            alw = al1 - 1.0;
            al2 = cluster_population(l2);
            alt = al2 + 1.0;

            cluster_center(1:dim_num,l1) = ...
              ( cluster_center(1:dim_num,l1) * al1 ...
              - point(1:dim_num,i) ) / alw;

            cluster_center(1:dim_num,l2) = ...
              ( cluster_center(1:dim_num,l2) * al2 ...
              + point(1:dim_num,i) ) / alt;

            cluster_population(l1) = cluster_population(l1) - 1;
            cluster_population(l2) = cluster_population(l2) + 1;
            an2(l1) = alw / al1;

            if ( 1.0 < alw )
              an1(l1) = alw / ( alw - 1.0 );
            else
              an1(l1) = Inf;
            end

            an1(l2) = alt / al2;
            an2(l2) = alt / ( alt + 1.0 );
            cluster(i) = l2;
            cluster2(i) = l1;

          end

        end

      end
%
%  If no re-allocation took place in the last POINT_NUM steps, return.
%
      if ( count == point_num )
        return
      end

    end

  end

  return
end


function [ crange ] = SpecColorRange( spec,range )
%[ crange ] = SpecColorRange( spec,range ) calculates a good color range for a
%spectrogram and sets the plot to the good range.
%
%
%DLevenstein 2017
%%
if ~exist('range','var')
    range = [2.5 2.5];
end

[~,mu,sig] = zscore(spec);
crange = [min(mu)-range(1)*max(sig) max(mu)+range(2)*max(sig)];
caxis(crange)

end


function value = p21_exact ( dim_num )

%*****************************************************************************80
%
%% P21_EXACT returns the exact integral for problem 21.
%
%  Licensing:
%
%    This code is distributed under the GNU LGPL license.
%
%  Modified:
%
%    19 January 2008
%
%  Author:
%
%    John Burkardt
%
%  Parameters:
%
%    Input, integer DIM_NUM, the spatial dimension.
%
%    Output, real VALUE, the exact value of the integral.
%
  c = 0.0;
  c = p21_r8 ( 'G', 'C', c );

  e = [];
  e = p21_i4vec ( 'G', 'E', dim_num, e );

  if ( any ( mod ( e(1:dim_num), 2 ) == 1 ) )
    value = 0.0;
    return
  end

  value = 2.0 * c;
  for i = 1 : dim_num
    arg = ( e(i) + 1 ) / 2.0;
    value = exact * gamma ( arg );
  end

  arg = ( sum ( e(1:dim_num) ) + dim_num ) / 2.0;
  value = value / gamma ( arg );

  return
end


clear all; close all; clc;
% Problem setup parameters
rng('default');
% Create the directory for storing results
[status_code,message,message_id] = mkdir('bin');

% We will load the details of this setup from a helper function
problem = barbara_problem();
num_training_patches = size(problem.training_patches, 2);
% we will use only 1000 training patches
chosen_training_patches = randperm(num_training_patches, 500);
% We will now train the dictionary
% Create the learning algorithm
trainer = spx.dictlearn.KSVD_OMP(problem.D, problem.N);
% We fix the sparsity level of representations
ksvd.K = 4;
trainer.InitialDictionary = problem.initial_dictionary;
% Option for not touching the  DC component atom
trainer.SpecialAtoms = 1;
% We want OMP to stop when residual is small
trainer.StopOnResidualNorm = true;
trainer.StopOnResNormStable = true;
%ksvd.MaxIterations = 10;
% Learn dictionary
trainer.train(problem.training_patches(:,chosen_training_patches));
% Get the learnt dictionary
learnt_dict = trainer.Dict;

save('bin/ex_barbara_ksvd.mat', 'learnt_dict');


function im_opp = rgb_to_opp(im)
% Transforms image into opponent color space
% Code by Ross Girshick

if ~isa(im, 'double')
  im = double(im);
end

im_opp = zeros(size(im));
% Color saliency boosting? c1*0.850, c2*0.524, c3*0.065
im_opp(:,:,1) = (im(:,:,1) - im(:,:,2)) / sqrt(2);
im_opp(:,:,2) = (im(:,:,1) + im(:,:,2) - 2*im(:,:,3)) / sqrt(6);
im_opp(:,:,3) = sum(im,3) / sqrt(3);
%min(im_opp(:)) % about -50
%max(im_opp(:)) % about 450
im_opp = max(0, min(255, im_opp)); % is this a good idea?


function [varargout] = binotest_dependent(X, Po)
% This function runs several different types of tests on dependent binomial data.
%
% :Usage:
% ::
%
%     [varargout] = binotest_dependent(X, Po)
%
% Overall, it tests the number of "hits" for each subject (row in X) against a null-hypothesis
% proportion p, across all subjects using a Z-test (two-tailed).
% The second level null hypothesis should be approximated by a normal
% distribution with a mean of p.  This approach assumes that each subject
% has an equal number of independent Bernoulli trials (columns in X) and
% that the number of subjects exceeds n=20 the test will be more accurate as n -> infinity.
% Also, calculates tests for each separate trial (e.g., subject columns),
% and the difference between proportions (two proportion z-test).
%
% :Inputs:
%
%   **X:**
%        X is a matrix of "hits" and "misses", coded as 1s and 0s.
%        where rows = subjects and columns = observations within
%        subject
%
%   **Po:**
%        Po is the null hypothesis proportion of "hits", e.g., often p = 0.5
%
% :Outputs:
%
%   **RES [1:5]:**
%        a structure containing the output of the stats for the
%        z-test, includes the number of subject (N), number of overall
%        hits (hits), the overall proportion of hits (prop), the standard
%        deviation (SE), z-statistic (Z), and the two tailed p-value (pval)
%        trial across all subjects.  Assumes independence
%
%   **RES1:**
%        Independent Single Interval Test for Column 1 (Column 1 only against Po)
%
%   **RES2:**
%        Independent Single Interval Test for Column 2 (Column 2 only against Po)
%
%   **RES3:**
%        Two proportion dependent difference z-test (Column 1 minus Column 2 against 0)
%
%   **RES4:**
%        Dependent single-interval test (Mean of Column 1 and Column 2 against Po)
%
%   **RES5:**
%        Two proportion dependent addition z-test (Column 1 plus Column 2 against 2 * Po)
%              (Similar to mean, not sure what this will be used for)
%
% :Examples:
% ::
%
%    [RES1, RES2, RES3, RES4, RES5] = binotest_dependent([1,1,1,1,0; 1,0,1,0,1]',.5)
%
% ..
%     Author and copyright information:
%
%     Copyright (C) 2014  Luke Chang & Tor Wager
%
%     This program is free software: you can redistribute it and/or modify
%     it under the terms of the GNU General Public License as published by
%     the Free Software Foundation, either version 3 of the License, or
%     (at your option) any later version.
%
%     This program is distributed in the hope that it will be useful,
%     but WITHOUT ANY WARRANTY; without even the implied warranty of
%     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
%     GNU General Public License for more details.
%
%     You should have received a copy of the GNU General Public License
%     along with this program.  If not, see <http://www.gnu.org/licenses/>.
% ..

% ..
%    Process Data
% ..

X = double(X); % just in case
[N, k] = size(X);  % number of subjects x dependent obs per person
A = X(:, 1);
B = X(:, 2);

% -------------------------------------------------------------------------
% Test if Data is in Correct Format
% -------------------------------------------------------------------------
if k > 2, error('NOT IMPLEMENTED FOR MORE THAN 2 DEP OBSERVATIONS PER UNIT'); end

u = unique(X(:));
if any(u ~= 1 & u ~= 0), error('X MUST HAVE VALUES OF 1 OR 0 ONLY.'); end

if N < 20
    warning(['Running dependent binomial test requires many subjects to approximate normal distribution.  You are only using ' num2str(length(unique(subject_id))) ' subjects.  Interpret results with caution.'])
end

% -------------------------------------------------------------------------
% Calculate Variance and Covariance
% -------------------------------------------------------------------------
% First, calculate variance of the quantity we want for a single trial,
% using Pa and Pb as estimates for the binomial parameters
% (This is not the SD of the SAMPLING distribution, but the SD of the
% distribution of interest (e.g., variance of sum, difference)
% Later, we will use this to get the SD for the SAMPLING distribution,
% which is used to construct confidence intervals and tests.

%Calculate Probabilities of each Subject Trial
Pa = sum(A) ./ N;  % P-hat for A
Pb = sum(B) ./ N;  % P-hat for B

% Calculate Variance for each Subject Trial
Va = Pa * (1 - Pa); % ./ N;
Vb = Pb * (1 - Pb); % ./ N;

% Calculate Covariance between Trials
% C = ( (sum(A .* B) ./ N) * (sum(~A .* ~B) ./ N) ) - ( (sum(A .* ~B) ./ N) * (sum(B .* ~A) ./ N) ); % % Agresti, p. 410 SAME as below
C = ( (sum(A .* B) ./ N) - (Pa * Pb) );   % Covariance - expected prop under independence

% Calculate Variance for Test Statistics adjusting for covariance
Vaplusb = Va + Vb - (2 * C);
Vameanb = ( Va + Vb - (2 * C) ) / 4;  % / 4 because V(aX) = a^2Var(X)
Vaminusb = Va + Vb + (2 * C);

% -------------------------------------------------------------------------
% Define Functions
% -------------------------------------------------------------------------

sdfunction = @(V, N) (V ./ N) .^ .5;  % Calculate Standard Deviation
zfunc = @(Px, Po, SD) (Px - Po) ./ SD; % Calculate Z-Value
pfunction_gauss = @(z) 2 .* min(normcdf(z, 0, 1), normcdf(z, 0, 1, 'upper')); %Calculate p-value for z-distribution

% pfunction_bino = @(N, Po, hits) 2 * min(binocdf(hits, N, Po), (1 - binocdf(hits - 1, N, Po))); %Calculate p-value for binomial distribution

% -------------------------------------------------------------------------
% Define Various Confidence Intervals - TO DO - See Agresti Book
% -------------------------------------------------------------------------

% switch  cimethod
%     case 'Wald'
%
%         cifunc = @(z)
%
%     case 'scoring'  % default
%
%     otherwise
%         error('Unknown CI method. Check inputs.');
% end

% Now, we will use this to get the SD for the SAMPLING distribution,
% which is used to construct confidence intervals and tests.
% e.g., SDa = (Va ./ N) .^ .5;

% -------------------------------------------------------------------------
% Calculate test statistics
% -------------------------------------------------------------------------
% Use z-test (large-sample approximation; ok for N = 20+)

names = {'P1' 'P2' 'P1minusP2' 'PavgP1P2' 'P1plusP2' };
est = [Pa Pb Pa-Pb mean([Pa Pb]) Pa+Pb];
Vs = [Va Vb Vaminusb Vameanb Vaplusb];
Pnull = [Po Po 0 Po 2*Po];

for i = 1:length(names)
    SDs(i) = sdfunction(Vs(i), N);
    Z(i) = zfunc(est(i), Pnull(i), SDs(i));
    pval(i) = pfunction_gauss(Z(i));
    %ci(i) = add confidence intervals

    RES{i} = struct('n', N, 'hits', est(i).*N, 'prop', est(i), 'SE', SDs(i), 'Z', Z(i), 'p_val', pval(i));

    varargout{i} = RES{i};
end

end % function


function test_conjgradmv()

N = 100;
x = rand(N,1);
A = sparse(double(rand(N) < 0.1));
K = A*A' + speye(N);

ind = 1:2:N;

xmv = x(ind);
Kmv = K(ind,ind);
ymv = Kmv*xmv;
y = zeros(N,1);
y(ind) = ymv;
I = false(N,1);
I(ind) = true;

x_pcg = pcg(Kmv,ymv,1e-6,100);
x_cg = conjgradmv(K,y,I,'tol',1e-6,'maxiter',100,'verbose',true);

norm(x_pcg-x_cg(ind))


function test95()
% function test95()
%
% Test for symfixedrankNewYYquotientfactory geometry (low-rank PSD matrix completion)
% This test is different from 'test98'. We use the tuned geometry,
% symfixedrankNewYYquotientfactory.
%

% This file is part of Manopt: www.manopt.org.
% Original author: Bamdev Mishra, Dec. 30, 2012.
% Contributors:
% Change log:

clear all; clc; close all;

m = 500;
r = 5;
A = randn(m, r);
C = A*A';

problem.M = symfixedrankNewYYquotientfactory(m, r);

df = problem.M.dim();
p = 5*df/(m*m);
symm = @(M) .5*(M+M');

% mask = symm(rand(m, m)) <= p;
mask = spones(sprandsym(m, p));

prop_known = sum(sum(mask == 1))/(m*m);

fprintf('Fraction of entries given: %f \n', full(prop_known));

problem.cost = @cost;
    function f = cost(X)
        f = 0.5*(norm(mask.*(X.Y*X.Y' - C), 'fro')^2);
    end

problem.grad = @grad;
    function g = grad(X)
        Y = X.Y;
        S = 2*mask.^2 .* (Y*Y' - C);
        r = size(Y, 2);
        YtY = Y'*Y;
        invYtY = eye(r) / YtY;
        g = struct('Y', S*Y*invYtY);
    end

problem.hess = @hess;
    function Hess = hess(X, eta)
        Y = X.Y;
        S = 2*mask.*( X.Y*X.Y' - C);
        S_star  = 2*mask.*(eta.Y*X.Y' + X.Y*eta.Y');

        r = size(Y, 2);
        YtY = Y'*Y;
        invYtY = eye(r) / YtY;

        Hess.Y = S_star*X.Y*invYtY;
        Hess.Y = Hess.Y + S*eta.Y*invYtY;
        Hess.Y = Hess.Y - 2*S*Y*(invYtY * symm(eta.Y'*X.Y) * invYtY);

        gradY = S*Y*invYtY;

        % I still need a correction factor for the non-constant metric
        Hess.Y = Hess.Y + gradY*symm(eta.Y'*X.Y)*invYtY + eta.Y*symm(gradY'*X.Y)*invYtY - X.Y*symm(eta.Y'*gradY)*invYtY;

        Hess = problem.M.proj(X, Hess);

    end

% % Check numerically whether gradient and Hessian are correct
%     checkgradient(problem);
%     drawnow;
%     pause;
%     checkhessian(problem);
%     drawnow;
%     pause;

% Initialization
[U, S, ~ ] = svds(mask.*C, r);
Y0 = U*(S.^0.5);
X0 = struct('Y', Y0);

% Options (not mandatory)
options.maxiter = inf;
options.maxinner = 30;
options.maxtime = 120;
options.tolgradnorm = 1e-9;
options.Delta_bar = m * r;
options.Delta0 = options.Delta_bar / 4;

% Pick an algorithm to solve the problem
[Xopt costopt info] = trustregions(problem, X0, options);
% [Xopt costopt info] = steepestdescent(problem, X0, options);
% [Xopt costopt info] = conjugategradient(problem, X0, options);

end


function test_poincareball()

    k = 3;
    n = 10;
    manifold = poincareballfactory(k, n);

    problem.M = manifold;

    P = manifold.rand();
    problem.cost = @(X) .5*norm(X-P, 'fro').^2;
    problem.egrad = @(X) X-P;
    problem.ehess = @(X, U) U;

    % Check numerically whether gradient and Ressian are correct
    checkgradient(problem);
    drawnow;
    pause;
    checkhessian(problem);
    drawnow;
    pause;

    % Initialization
    X0 = [];

    % Options (not mandatory)
    options.maxiter = inf;
    options.maxinner = 30;
    options.maxtime = 120;
    options.tolgradnorm = 1e-10;

    % Pick an algorithm to solve the problem
    Xopt = trustregions(problem, X0, options);
    % [Xopt, costopt, info] = steepestdescent(problem, X0, options);
    % [Xopt, costopt, info] = conjugategradient(problem, X0, options);

    % Curiously enough, the optimizer sometimes doesn't find the solution
    % P. When this happens, the erroneous columns of Xopt tend to have norm
    % very close to 1, which is also where numerics become tricky.
    Xopt - P
    sum(Xopt.^2)
    sum(P.^2)

    evs = hessianspectrum(problem, Xopt);
    evs = real(evs);
    stairs(sort(evs));
    title(['Eigenvalues of the Hessian of the cost function ' ...
           'at the solution']);
    fprintf('Hessian condition number at solution: %g\n', ...
            max(abs(evs))/min(abs(evs)));
end


function f = flops(fl)
% FLOPS         Get or set the global flopcount variable.
% FLOPS returns the current flopcount.
% FLOPS(F) sets flopcount to F.
%
% 0 flops: -x ' repmat
% 1 flop each: + - .*
% 2 flops each: < > == ~=
% For complex numbers, + is 2 flops, * is 6 flops.
% col_sum(x) takes (rows(x)-1)*cols(x) flops (use FLOPS_COL_SUM).
% row_sum(x) takes rows(x)*(cols(x)-1) flops (use FLOPS_ROW_SUM).
% Use FLOPS_DIV for ./
% Use FLOPS_RANDNORM for randn
% Use FLOPS_SQRT for sqrt
% Use FLOPS_POW for .^
% Use FLOPS_EXP for exp
% Use FLOPS_LOG for log, sin, and other special functions.
%
% See FLOPS_MUL, FLOPS_SOLVE, FLOPS_INV, FLOPS_CHOL, FLOPS_DET, ...

global flopcount;
if nargin == 1
  flopcount = fl;
  if nargout == 1
    f = fl;
  end
else
  f = flopcount;
end


function [Q, qMap] = mef_ssim(imgSeq, fI,  K, window)

if (nargin < 2 || nargin > 4)
   Q = -Inf;
   qMap = Inf;
   return;
end

if (~exist('K', 'var'))
   K = 0.03;
end

if (~exist('window', 'var'))
   window = fspecial('gaussian', 11, 1.5);
end

imgSeq = double(imgSeq);
fI = double(fI);
[s1, s2, s3] = size(imgSeq);
wSize = size(window,1);
sWindow = ones(wSize) / wSize^2; % square window used to calculate the distance
bd = floor(wSize/2);
mu = zeros(s1-2*bd, s2-2*bd, s3);
ed = zeros(s1-2*bd, s2-2*bd, s3);
for i = 1:s3
    img = squeeze(imgSeq(:,:,i));
    mu(:,:,i) = filter2(sWindow, img, 'valid');
    muSq = mu(:,:,i) .* mu(:,:,i);
    sigmaSq = filter2(sWindow, img.*img, 'valid') - muSq;
    ed(:,:,i) =  sqrt( max( wSize^2 * sigmaSq, 0 ) ) + 0.001; % add a small constant to avoid instability
end

R = zeros(s1-2*bd,s2-2*bd); % consistency map which could be used as an output if necessary
for i = bd+1:s1-bd
    for j = bd+1:s2-bd
        vecs = reshape(imgSeq(i-bd:i+bd,j-bd:j+bd,:),[wSize*wSize, s3]);
        denominator = 0;
        for k = 1:s3
            denominator = denominator + norm(vecs(:,k) - mu(i-bd,j-bd,k));
        end
        numerator = norm(sum(vecs,2) - mean(sum(vecs,2)));
        R(i-bd,j-bd) = (numerator + eps) / (denominator + eps);
    end
end

R(R > 1) = 1 - eps; % get rid of numerical instability
R(R < 0) = 0 + eps;

p = tan(pi/2 * R);
p( p >  10 ) = 10; % to avoid blow up (large number such as 10 is equivalent to taking maximum)
p = repmat(p,[1,1,s3]);

wMap = (ed / wSize).^p + eps; % to avoid blowing up
normalizer = sum(wMap,3);
wMap = wMap ./ repmat(normalizer,[1,1,s3]);
maxEd = max(ed,[],3);

C = (K * 255)^2;
qMap = zeros(s1-2*bd, s2-2*bd);
for i = bd+1:s1-bd
    for j = bd+1:s2-bd
        blocks = imgSeq(i-bd:i+bd,j-bd:j+bd,:);
        rBlock = zeros(wSize,wSize);
        for k = 1 : s3
            rBlock = rBlock  + wMap(i-bd,j-bd,k) * ( blocks(:,:,k) - mu(i-bd,j-bd,k) ) / ed(i-bd,j-bd,k);
        end
        if norm(rBlock(:)) > 0
            rBlock = rBlock / norm(rBlock(:)) * maxEd(i-bd,j-bd);
        end
        fBlock = fI(i-bd:i+bd,j-bd:j+bd);
        rVec = rBlock(:);
        fVec = fBlock(:);
        mu1 = sum( window(:) .* rVec );
        mu2 = sum( window(:) .* fVec );
        sigma1Sq = sum( window(:) .* (rVec - mu1).^2 );
        sigma2Sq = sum( window(:) .* (fVec - mu2).^2 );
        sigma12 = sum(  window(:) .* (rVec - mu1) .* (fVec - mu2)  );
        qMap(i-bd,j-bd) = ( 2 * sigma12 + C ) ./ ( sigma1Sq + sigma2Sq + C );
    end
end

Q = mean2(qMap);


clear, clc;

% This is an example for running the function epp
%
%  Problem:
%
%  min  1/2 || x - y||^2 + rho * \|x\|_q
%
%  which is a subproblem of the Euclidean projection:
%
%  min \sum_i 1/2 || X_i - Y_i||^2 + rho * \sum_i \|X_i\|_q
%
% For detailed description of the function, please refer to the Manual.
%
%% ------------   History --------------------
% First version on August 10, 2009.
%
%
% For any problem, please contact Jun Liu (j.liu@asu.edu)

cd ..
cd ..

root=cd;
addpath(genpath([root '/SLEP']));
                     % add the functions in the folder SLEP to the path

% change to the original folder
cd Examples/L1Lq;

n=100;
v=randn(n, 1);
lambda=1;
q=3;
c0=0.1;

% run the function epp
[x, c, iter_step]=epp(v, n, lambda, q, c0);


%% test modulate for all oasis functions.
col = {[0 114 178],[0 158 115], [213 94 0],[230 159 0],...
    [86 180 233], [204 121 167], [64 224 208], [240 228 66]}; % colors
plot_cvx = false;

%% example 2: foopsi, AR2 model
g = [1.7, -0.712];         % AR coefficient
noise = 1;
T = 3000;
framerate = 30;
firerate = 0.5;
b = 0;              % baseline
N = 20;              % number of trials
seed = 3;          % seed for genrating random variables
[Y, trueC, trueS] = gen_data(g, noise, T, framerate, firerate, b, N, seed);
y = Y(1,:);
true_c = trueC(1,:);  %#ok<*NASGU>
true_s = trueS(1,:);
% case 1: all parameters are known
lambda = 25;
[c_oasis, s_oasis] = deconvolveCa(y, 'ar2', g, 'foopsi', 'lambda', lambda);  %#ok<*ASGLU>

figure('name', 'FOOPSI, AR2, known: g, lambda', 'papersize', [15, 4]);
show_results;

% case 2: know lambda
lambda = 2.5;
[c_oasis, s_oasis, options] = deconvolveCa(y, 'ar2', 'sn', noise, 'foopsi', 'lambda',...
    lambda);
fprintf('true gamma:        %.3f\t %.3f\n', g(1), g(2));
fprintf('estimated gamma:   %.3f\t %.3f\n', options.pars(1),  options.pars(2));

figure('name', 'FOOPSI, AR2, known:lambda, estimated: g', 'papersize', [15, 4]);
show_results;

%%%%%%%%%%%%%%  END %%%%%%%%%%%%%%%%%%


% [] = XXXX()
%
% -----------------------------Definition---------------------------------%
%
% description
%
% usage :
%
% -----------------------------Input--------------------------------------%
%
% XX : description + dimension
%
% -----------------------------Output-------------------------------------%
%
% XX : description + dimension
%
% -----------------------------References---------------------------------%
%
% [1] : XXX
%
%
%   Project : BCI-EEG
%
%   author : A. Barachant
%   date : 2011-XXXX
%   version : 1.0
%   status : a terminer
%   CEA/GRENOBLE-LETI/DTBS
%
%   See also MEAN, MEDIAN, RIEMANN_MEAN, REIMANN_MEDIAN, RIEMANN_TRIMMED_MEAN.

% [EOF: XXX.m]

function C = mean_covariances(COV,method_mean,arg_mean)

if (nargin<2)||(isempty(method_mean))
    method_mean = 'arithmetic';
end
if (nargin<3)
    arg_mean = {};
end

switch method_mean
    case 'riemann'
        C = riemann_mean(COV,arg_mean);
    case 'riemanndiag'
        C = riemann_diag_mean(COV,arg_mean);
    case 'riemanntrim'
        C = riemann_trimmed_mean(COV,arg_mean);
    case 'median'
        C = median(COV,3);
    case 'riemannmed'
        C = riemann_median(COV,arg_mean);
    case 'logeuclid'
        C = logeuclid_mean(COV);
    case 'opttransp'
        C = opttransp_mean(COV);
    case 'ld'
        C = logdet_mean(COV);
    case 'geodesic'
        C = geodesic_mean(COV,arg_mean);
    case 'harmonic'
        iCOV = zeros(size(COV));
        for i=1:size(COV,3)
            iCOV(:,:,i) = inv(COV(:,:,i));
        end
        C = inv(mean(iCOV,3));
    case 'geometric'
        B = COV(:,:,1);
        for i=2:size(COV,3)
            B = B*COV(:,:,i);
        end
        C = B^(1/size(COV,3));
    otherwise
        C = mean(COV,3);
end


% rigid registration of frames with offsets ds
function dreg = register_movie(data, ops, ds)

orig_class = class(data);

% compute maximum batch size for GPU
[h, w, nFrames] = size(data);
if getOr(ops, 'useGPU', 0)
    nFramesPerBatch = getBatchSize(w*h);
else
    nFramesPerBatch = 1000; % if not on GPU, should have plenty of RAM available
end
%%
nFramesPerBatch = 8;

nBatches = ceil(nFrames/nFramesPerBatch);
startFrame = 1:nFramesPerBatch:nFrames;
endFrame = min(startFrame+nFramesPerBatch-1, nFrames);
dreg = zeros(size(data), orig_class);

for iBatch = 1:nBatches
    idx = startFrame(iBatch):endFrame(iBatch);
    if ops.useGPU
        dataBatch = gpuArray(single(data(:,:,idx)));
    else
        dataBatch = data(:,:,idx);
    end
    [Ly, Lx, NT] = size(dataBatch);

    Ny = ifftshift([-fix(Ly/2):ceil(Ly/2)-1]);
    Nx = ifftshift([-fix(Lx/2):ceil(Lx/2)-1]);
    [Nx,Ny] = meshgrid(Nx,Ny);
    Nx = Nx / Lx;
    Ny = Ny / Ly;

    if ops.useGPU
        dsBatch = gpuArray(permute(ds(idx, :), [3, 2, 1]));
        Nx = gpuArray(single(Nx));
        Ny = gpuArray(single(Ny));
    else
        dsBatch = ds(idx, :);
    end

    if ops.useGPU % do it batch-by-batch
        dph         = 2*pi*(bsxfun(@times, dsBatch(1,1,:), Ny) + ...
            bsxfun(@times, dsBatch(:,2,:), Nx));
        fdata       = fft2(dataBatch);
        dregBatch   = gather_try(real(ifft2(fdata .* exp(1i * dph))));
    else % do it frame-by-frame
        dregBatch = zeros(size(dataBatch), orig_class);
        for i = 1:NT
            dph         = 2*pi*(dsBatch(i,1)*Ny + dsBatch(i,2)*Nx);
            fdata       = fft2(single(dataBatch(:,:,i)));
            dregBatch(:,:,i) = real(ifft2(fdata .* exp(1i * dph)));
        end
    end

    dreg(:,:,idx) = dregBatch;
end


function [mg,nu,sig,info] = spm_rice_mixture(h,x,K)
% Fit a mixture of Ricians to a histogram
% FORMAT [mg,nu,sig] = spm_rice_mixture(h,x,K)
% h    - histogram counts
% x    - bin positions (plot(x,h) to see the histogram)
% K    - number of Ricians
% mg   - integral under each Rician
% nu   - "mean" parameter of each Rician
% sig  - "standard deviation" parameter of each Rician
% info - This struct can be used for plotting the fit as:
%            plot(info.x(:),info.p,'--',info.x(:), ...
%                 info.h/sum(info.h)/info.md,'b.', ...
%                 info.x(:),info.sp,'r');
%
% An EM algorithm is used, which involves alternating between computing
% belonging probabilities, and then the parameters of the Ricians.
% The Koay inversion technique is used to compute the Rician parameters
% from the sample means and standard deviations. This is described at
% https://en.wikipedia.org/wiki/Rician_distribution
%__________________________________________________________________________
% Copyright (C) 2012-2019 Wellcome Trust Centre for Neuroimaging

% John Ashburner
% $Id: spm_rice_mixture.m 7595 2019-05-23 13:48:53Z mikael $

mg  = ones(K,1)/K;
nu  = (0:(K-1))'*max(x)/(K+1);
sig = ones(K,1)*max(x)/K/10;
lam = (sum(x.*h)/sum(h)/K).^2;

m0 = zeros(K,1);
m1 = zeros(K,1);
m2 = zeros(K,1);
ll = -Inf;
for iter=1:10000
    p  = zeros(numel(x),K);
    for k=1:K
        % Product Rule
        % p(class=k, x | mg, nu, sig) = p(class=k|mg) p(x | nu, sig, class=k)
        p(:,k) = mg(k)*ricepdf(x(:),nu(k),sig(k)^2) + eps;
    end

    % Sum Rule
    % p(x | mg, nu, sig) = \sum_k p(class=k, x | mg, nu, sig)
    sp  = sum(p,2);
    oll = ll;
    ll  = sum(log(sp).*h(:)); % Log-likelihood
    if ll-oll<1e-8*sum(h), break; end

    %fprintf('%g\n',ll);
    %md = mean(diff(x));
    %plot(x(:),p,'--',x(:),h/sum(h)/md,'b.',x(:),sp,'r'); drawnow

    % Bayes Rule
    % p(class=k | x, mg, nu, sig) = p(class=k, x | mg, nu, sig) / p(x | mg, nu, sig)
    p = bsxfun(@rdivide,p,sp);

    % Compute moments from the histograms, weighted by the responsibilities (p).
    for k=1:K
        m0(k) = sum(p(:,k).*h(:));              % Number of voxels in class k
        m1(k) = sum(p(:,k).*h(:).*x(:));        % Sum of the intensities in class k
        m2(k) = sum(p(:,k).*h(:).*x(:).*x(:));  % Sum of squares of intensities in class k
    end

    mg = m0/sum(m0); % Mixing proportions
    for k=1:K
        mu1 = m1(k)./m0(k);                                % Mean
        mu2 = (m2(k)-m1(k)*m1(k)/m0(k)+lam*1e-3)/(m0(k)+1e-3); % Variance

        % Compute nu & sig from mean and variance
        [nu(k),sig(k)] = moments2param(mu1,mu2);
    end
    %disp([nu'; sig'])
end

if nargout >= 4
    % This info can be used for plotting the fit
    info    = struct;
    info.x  = x;
    info.h  = h;
    info.p  = p;
    info.sp = sp;
    info.md = mean(diff(x));
end
%__________________________________________________________________________

%__________________________________________________________________________
function [nu,sig] = moments2param(mu1,mu2)
% Rician parameter estimation (nu & sig) from mean (mu1) and variance
% (mu2) via the Koay inversion technique.
% This follows the scheme at
% https://en.wikipedia.org/wiki/Rice_distribution#Parameter_estimation_.28the_Koay_inversion_technique.29
% This Wikipedia description is based on:
% Koay, C.G. and Basser, P. J., Analytically exact correction scheme
% for signal extraction from noisy magnitude MR signals,
% Journal of Magnetic Resonance, Volume 179, Issue = 2, p. 317–322, (2006)

r     = mu1/sqrt(mu2);
theta = sqrt(pi/(4-pi));
if r>theta
    for i=1:256
        xi    = 2+theta^2-pi/8*exp(-theta^2/2)*((2+theta^2)*besseli(0,theta^2/4)+theta^2*besseli(1,theta^2/4))^2;
        g     = sqrt(xi*(1+r^2)-2);
        if abs(theta-g)<1e-6, break; end
        theta = g;
    end
    if ~isfinite(xi), xi = 1; end
    sig = sqrt(mu2)/sqrt(xi);
    nu  = sqrt(mu1^2+(xi-2)*sig^2);
else
    nu  = 0;
    sig = (2^(1/2)*(mu1^2 + mu2)^(1/2))/2;
end
%__________________________________________________________________________

%__________________________________________________________________________
function p = ricepdf(x,nu,sig2)
% Rician PDF
% p = ricepdf(x,nu,sig2)
% https://en.wikipedia.org/wiki/Rice_distribution#Characterization
p       = zeros(size(x));
tmp     = -(x.^2+nu.^2)./(2*sig2);
msk     = (tmp > -95) & (x*(nu/sig2) < 85) ; % Identify where Rice probability can be computed
p(msk)  = (x(msk)./sig2).*exp(tmp(msk)).*besseli(0,x(msk)*(nu/sig2)); % Use Rician distribution
p(~msk) = (1./sqrt(2*pi*sig2))*exp((-0.5/sig2)*(x(~msk)-nu).^2);      % Use Gaussian distribution


function dista=w_distance(p1, p2)
    vec = p1-p2;
    dista=sqrt(sum(vec.^2));
end


function f = dtlz7a(x)

%% cost functions
f(:,1) = x(:,1);
f(:,2) = x(:,2);

%% g function
sum1 = 0;
sum2 = 0;
for i = 3:8
sum1 = sum1 + x(:,i);
end
g = 1+(9/6)*sum1;
for i = 1:2
sum2 = sum2 + (f(:,i)/(1+g))*(1+sin(3*pi*f(:,i)));
end
h = 3 - sum2;

%% cost functions
f(:,3) = (1+g)*h;


function x=nsst_rec1(dst,lpfilt)
% This function performs the inverse (local) nonsubsampled shearlet transform as given
% in G. Easley, D. Labate and W. Lim, "Sparse Directional Image Representations
% using the Discrete Shearlet Transform", Appl. Comput. Harmon. Anal. 25 pp.
% 25-46, (2008).
%
% Input:
%
% dst          - the nonsubsampled shearlet coefficients
%
% lpfilt       - the filter to be used for the Laplacian
%                Pyramid/ATrous decomposition using the codes
%                written by Arthur L. Cunha
%
% Output
%
% x         - the reconstructed image
%
% Code contributors: Glenn R. Easley, Demetrio Labate, and Wang-Q Lim.
% Copyright 2011 by Glenn R. Easley. All Rights Reserved.
%

level=length(dst)-1;
y{1}=dst{1};
for i=1:level,
      y{i+1} = real(sum(dst{i+1},3));
end

x=real(atrousrec(y,lpfilt));


function window_sz = get_search_window( target_sz, im_sz)
% GET_SEARCH_WINDOW

% if(target_sz(1)/target_sz(2) > 2)
%     % For objects with large height, we restrict the search window with padding.height
%     window_sz = floor(target_sz.*[1+padding.height, 1+padding.generic]);
%
% elseif(prod(target_sz)/prod(im_sz(1:2)) > 0.05)
%     % For objects with large height and width and accounting for at least 10 percent of the whole image,
%     % we only search 2x height and width
%     window_sz=floor(target_sz*(1+padding.large));
%
% else
%     %otherwise, we use the padding configuration
%     window_sz = floor(target_sz * (1 + padding.generic));

ratio=target_sz(1)/target_sz(2);
if ratio>1
    window_sz=round(target_sz.*[5,5*ratio]);
else
    window_sz=round(target_sz.*[5/ratio,5]);
end

%window_sz=round(target_sz.*[5,9]);
window_sz=window_sz-mod(window_sz,2)+1;

end


% Kbeta
%
% Computes weighted sum of kernel matrices
%
% Usage:
%	K = Kbeta(Ks,w);
%	K = Kbeta(Ks,w,symmetric);
%
% Input:
%	Ks - Gram matrices [n x n x M]
%	w - weights [M x 1]
%	symmetric (=false) - optional argument. If 1 the Ks
%		matrices are assumed to be symmetric which results in a
%		small speedup.
%
% Output:
%	K - weighted kernel matrix equiv with
%		K=0;for m=1:M, K=K+ w(m)*Ks(:,:,m);end
%
% Example:
%
%	K = Kbeta(randn(100,100,10),rand(10,1));
%
% Peter Gehler 07/2008 pgehler@tuebingen.mpg.de


function x = r83_cr_sl ( n, a_cr, b )

%*****************************************************************************80
%
%% R83_CR_SL solves a real linear system factored by R83_CR_FA.
%
%  Discussion:
%
%    The matrix A must be tridiagonal.  R83_CR_FA is called to compute the
%    LU factors of A.  It does so using a form of cyclic reduction.  If
%    the factors computed by R83_CR_FA are passed to R83_CR_SL, then one or many
%    linear systems involving the matrix A may be solved.
%
%    Note that R83_CR_FA does not perform pivoting, and so the solution
%    produced by R83_CR_SL may be less accurate than a solution produced
%    by a standard Gauss algorithm.  However, such problems can be
%    guaranteed not to occur if the matrix A is strictly diagonally
%    dominant, that is, if the absolute value of the diagonal coefficient
%    is greater than the sum of the absolute values of the two off diagonal
%    coefficients, for each row of the matrix.
%
%  Example:
%
%    Here is how a R83 matrix of order 5 would be stored:
%
%       *  A12 A23 A34 A45
%      A11 A22 A33 A44 A55
%      A21 A32 A43 A54  *
%
%  Licensing:
%
%    This code is distributed under the GNU LGPL license.
%
%  Modified:
%
%    22 March 2004
%
%  Author:
%
%    MATLAB version by John Burkardt.
%
%  Reference:
%
%    Roger Hockney,
%    A fast direct solution of Poisson's equation using Fourier Analysis,
%    Journal of the ACM,
%    Volume 12, Number 1, pages 95-113, January 1965.
%
%  Parameters:
%
%    Input, integer N, the order of the matrix.
%    N must be positive.
%
%    Input, real A_CR(3,2*N+1), factorization information computed by R83_CR_FA.
%
%    Input, real B(N), the right hand side vector.
%
%    Output, real X(N), the solution of the linear system.
%
  if ( n <= 0 )
    fprintf ( 1, '\n' );
    fprintf ( 1, 'R83_CR_SL - Fatal error!\n' );
    fprintf ( 1, '  Nonpositive N = %d\n', n );
    return
  end

  if ( n == 1 )
    x(1) = a_cr(2,2) * b(1);
    return
  end
%
%  Set up RHS.
%
  rhs(1) = 0.0;
  rhs(2:n+1) = b(1:n);
  rhs(n+2:2*n+1) = 0.0;

  il = n;
  ndiv = 1;
  ipntp = 0;

  while ( 1 < il )

    ipnt = ipntp;
    ipntp = ipntp + il;
    il = floor ( il / 2 );
    ndiv = ndiv * 2;
    ihaf = ipntp;

    for iful = ipnt+2 : 2 : ipntp
      ihaf = ihaf + 1;
      rhs(ihaf+1) = rhs(iful+1) - a_cr(3,iful) * rhs(iful) ...
        - a_cr(1,iful+1) * rhs(iful+2);
    end

  end

  rhs(ihaf+1) = rhs(ihaf+1) * a_cr(2,ihaf+1);
  ipnt = ipntp;

  while ( 0 < ipnt )

    ipntp = ipnt;
    ndiv = floor ( ndiv / 2 );
    il = floor ( n / ndiv );
    ipnt = ipnt - il;
    ihaf = ipntp;

    for ifulm = ipnt+1 : 2 : ipntp
      iful = ifulm + 1;
      ihaf = ihaf + 1;
      rhs(iful+1) = rhs(ihaf+1);
      rhs(ifulm+1) = a_cr(2,ifulm+1) * ( rhs(ifulm+1) - a_cr(3,ifulm) * rhs(ifulm) ...
        - a_cr(1,ifulm+1) * rhs(iful+1) );
    end

  end

  x(1:n) = rhs(2:n+1);

  return
end


% Compute geodesic distance along the mesh
%[V,F] = readOBJ('data/spot_good.obj'); % this one will work
[V,F] = readOBJ('data/spot_bad.obj');  % but this one will not!

% NOTE: The result may not be the same on the "bad" mesh after you repair
% it, because this code computes the distance from vertex "1", but which
% vertex is vertex "1" may change. That is fine, as long as it outputs
% reasonable a looking distance.

% Run the algorithm
% (I promise nothing is wrong with this algorithm, the problem is the mesh)
dists = heat_geodesic(V,F,1);

% Plot the resulting function
t = tsurf(F,V, 'CData',dists);
shading interp;
axis equal;
axis off;
colormap(cbrewer('Blues', 500));
light('Position',[-1.5 1 1],'Style','local');
lights = camlight;
set(t, 'FaceLighting','gouraud', 'FaceColor','interp');
set(t, 'DiffuseStrength',0.5, 'SpecularStrength',0.2, 'AmbientStrength',0.3);
camproj('perspective');


function demo_gpfa(seed)

if nargin >= 1
  randn('state', seed);
  rand('state', seed);
end

D = 4;

%
% Generate some data
%

%
% Generate X
%

covfunc_x = cell(D,1);

N_x = 300;
in_x = 1:N_x;

X = zeros(D,N_x);
D2_xx = sq_dist(in_x);
D_xx = sqrt(D2_xx);

pseudo_x = in_x(:,1:8:end);
D2_pp = sq_dist(pseudo_x);
D2_px = sq_dist(pseudo_x, in_x);
d2_x = diag(D2_xx);

is_pseudos_x = false(D,1);

% Trend component
covfunc = @(D2) gp_cov_se(D2);
covfunc_x{1} = gp_cov_pseudo(gp_cov_jitter(covfunc(D2_pp), 1e-3), ...
                             covfunc(D2_px), ...
                             covfunc(d2_x));
theta_x{1} = [100]; % true length scale
theta_init_x{1} = [50]; % init length scale
is_pseudos_x(1) = true;

% Almost periodic component (with fixed wave length)
covfunc = @(D2) gp_cov_product(gp_cov_periodic(sqrt(D2), ...
                                               'wavelength', 30), ...
                               gp_cov_se(D2));
covfunc_x{2} = gp_cov_pseudo(gp_cov_jitter(covfunc(D2_pp), 1e-3), ...
                             covfunc(D2_px), ...
                             covfunc(d2_x));
theta_x{2} = [1  % smoothness
              400]; % length scale of the decay
theta_init_x{2} = [0.5  % smoothness
                   800]; % length scale of the decay
is_pseudos_x(2) = true;

% Slow component
covfunc_x{3} = gp_cov_jitter(gp_cov_rq(D2_xx));
theta_x{3} = [5   % length scale
              1]; % alpha ("degrees of freedom")
theta_init_x{3} = [5     % length scale
                   1]; % alpha ("degrees of freedom")
is_pseudos_x(3) = false;

% Fast component
covfunc_x{4} = gp_cov_jitter(gp_cov_pp(D_xx, 1));
theta_x{4} = [4]; % length scale
theta_init_x{4} = [2.5]; % length scale
is_pseudos_x(4) = false;

% Generate X by random samples
for d=1:D
  if ~is_pseudos_x(d)
    K = covfunc_x{d}(theta_x{d});
    L = chol(K, 'lower');
    X(d,:) = L*randn(N_x,1);
  else
    [K_pp, K_px, k_x] = covfunc_x{d}(theta_x{d});
    L = chol(K_pp, 'lower');
    X(d,:) = K_px'*linsolve_lchol(L, L*randn(size(L,1),1));
  end
end

%
% Generate W
%

N_w = 30;
in_w = rand(2,N_w);

D2_ww = sq_dist(in_w);
W = zeros(D,N_w);

% Meshgrid for showing the spatial functions
[grid_w1,grid_w2] = meshgrid(0:0.03:1, 0:0.03:1);
grid_w = [grid_w1(:)'; grid_w2(:)'];
D2_grid = sq_dist([grid_w, in_w]);
N_grid = size(grid_w, 2);
W_grid = zeros(D,N_grid);

% RQ covariance function for all spatial components
covfunc_w = cell(D,1);
covfunc_grid = cell(D,1);
covfunc_w(:) = {gp_cov_scale(gp_cov_jitter(gp_cov_rq(D2_ww)))};
covfunc_grid(:) = {gp_cov_scale(gp_cov_jitter(gp_cov_rq(D2_grid)))};
% True parameter values
theta_w{1} = [2     % signal magnitude
              0.4   % length scale
              3.0]; % alpha ("degrees of freedom")
theta_w{2} = [1     % signal magnitude
              0.2  % length scale
              2.0]; % alpha ("degrees of freedom")
theta_w{3} = [1     % signal magnitude
              0.1  % length scale
              1.0]; % alpha ("degrees of freedom")
theta_w{4} = [0.5   % signal magnitude
              0.05  % length scale
              1.0]; % alpha ("degrees of freedom")
% Parameter value initializations
theta_init_w = cell(D,1);
theta_init_w(:) = columns_to_cells(...
    [ones(1,D)            % signal magnitude
     linspace(0.3,0.01,D) % length scale
     ones(1,D)]);         % alpha ("degrees of freedom")
% $$$ theta_init_w(:) = {[1    % signal magnitude
% $$$                     0.02 % length scale
% $$$                     1]}; % alpha ("degrees of freedom")

% Generate W by random samples
for d=1:D
  K = covfunc_grid{d}(theta_w{d});
  L = chol(K, 'lower');
  w = L*randn(length(L),1);
  W_grid(d,:) = w(1:N_grid);
  W(d,:) = w((N_grid+1):end);
end

%
% Generate Y
%

% Noisy observations
Y_noiseless = W'*X;
s = 1;
Y = Y_noiseless + s*randn(N_w,N_x);

% Missing values
Y(rand(size(Y))<0.0) = nan;

%
% Plot
%

% True spatial components
figure
for d=1:D
  subplot(ceil(sqrt(D)), ceil(sqrt(D)), d);
  % Plot contours of W
  cmax = max(abs(W_grid(d,:)));
  contourf(grid_w1, grid_w2, reshape(W_grid(d,:), size(grid_w1)));
  set(gca, 'clim', [-cmax, cmax])
  map_colormap();
  % Plot locations
  hold on
  plot(in_w(1,:), in_w(2,:), 'k+');
end

% True temporal components
tsplot(X);

%
% Inference with GPFA
%

% GP module with component-wise factorization for X
%theta_init_x = theta_x; % initialize with true parameter values
X_module = factor_module_gp_factorized(N_x, covfunc_x, theta_init_x, ...
                                       'update_hyperparameters', [5 10:10:100], ...
                                       'maxiter_hyperparameters', 10, ...
                                       'is_pseudo', is_pseudos_x, ...
                                       'init', zeros(D,N_x));

% GP module with component-wise factorization for W
%theta_init_w = theta_w; % initialize with true parameter values
W_module = factor_module_gp_factorized(N_w, covfunc_w, theta_init_w, ...
                                       'update_hyperparameters', [5 10:10:100], ...
                                       'maxiter_hyperparameters', 10);

% Isotropic noise
noise_module = noise_module_isotropic(N_w, N_x, ...
                                      'prior', struct('a_tau', 1e-3, ...
				                      'b_tau', 1e-3), ...
                                      'init', struct('a_tau', 10, ...
				                     'b_tau', 1));

% Run GPFA
Q = gpfa(D, Y, W_module, X_module, noise_module, ...
         'maxiter', 50, ...
         'debug', false, ...
...%         'rotate', 1:10, ...
         'rotate', 1:15, ...
         'rotation_checkgrad', false, ...
         'rotation_show', false, ...
         'rotation_maxiter', 30);
% 'predict', true

fprintf('RMSE of the noiseless reconstruction: %f\n',  rmse(Y_noiseless, Q.W'*Q.X));
fprintf('STD of the noiseless predictions: %f\n', ...
        sqrt(mean(mean(Q.W.^2'*Q.CovX + Q.CovW'*Q.X.^2 + Q.CovW'*Q.CovX))));
fprintf('Estimated noise STD: %f\n', Q.Tau(1).^(-0.5))

% The development of the loglikelihood lower bound during the iteration
figure
plot(Q.loglikelihood)

% Estimated latent temporal signals
figure
for d=1:D
  subplot(D,1,d);
  errorplot(Q.X(d,:), 2*sqrt(Q.CovX(d,:)))
end

return

% $$$
% $$$
% $$$
% $$$
% $$$
% $$$
% $$$
% $$$
% $$$
% $$$ % $$$ tsplot(X)
% $$$ % $$$ tsplot(W)
% $$$ % $$$ tsplot(Y, 'k+')
% $$$ % $$$ return
% $$$
% $$$ % $$$ regularize = @(covfunc,N) gp_cov_sum(covfunc, ...
% $$$ % $$$                                      gp_cov_scale(gp_cov_delta(N), ...
% $$$ % $$$                                                   'scale', 1e-6));
% $$$
% $$$ %
% $$$ % GP module for X
% $$$ %
% $$$
% $$$ N_p = 2;
% $$$ pseudo_x = linspace(min(in_x),max(in_x), N_p);
% $$$ is_pseudo_x = false(D,1);
% $$$
% $$$ D2_xx = sq_dist(in_x);
% $$$ D2_pp = sq_dist(pseudo_x);
% $$$ D2_xp = sq_dist(in_x, pseudo_x);
% $$$
% $$$ % Covariance functions
% $$$ covfunc = @(D2) gp_cov_se(D2);
% $$$ covfunc_x = cell(D,1);
% $$$ covfunc_x{1} = gp_cov_jitter(gp_cov_se(D2_xx));
% $$$ % $$$ covfunc_x{1} = gp_cov_pseudo(gp_cov_jitter(gp_cov_se(D2_pp)), ...
% $$$ % $$$                              gp_cov_se(D2_xp), ...
% $$$ % $$$                              gp_cov_se(diag(D2_xx)));
% $$$ % $$$ is_pseudo_x(1) = true
% $$$ theta_x{1} = 30;
% $$$
% $$$ covfunc_x{2} = gp_cov_jitter(covfunc(D2_xx));
% $$$ theta_x{2} = 3;
% $$$ covfunc_x{3} = gp_cov_jitter(gp_cov_pp(sqrt(D2_xx),1));
% $$$ theta_x{3} = 1;
% $$$
% $$$ % $$$ figure
% $$$ % $$$ imagesc(covfunc_x{1}(theta_x{1}));
% $$$ % $$$ figure
% $$$ % $$$ imagesc(covfunc_x{2}(theta_x{2}));
% $$$ % $$$ return
% $$$
% $$$ % GP module with component-wise factorization
% $$$ X_module = factor_module_gp_factorized(N_x, ...
% $$$                                        covfunc_x, ...
% $$$                                        theta_x, ...
% $$$                                        'is_pseudo', is_pseudo_x, ...
% $$$                                        'update_hyperparameters', 2);
% $$$
% $$$ %
% $$$ % GP module for W
% $$$ %
% $$$
% $$$ D2_ww = sq_dist(in_w);
% $$$
% $$$ % Covariance functions
% $$$ covfunc_w = cell(D,1);
% $$$ covfunc_w{1} = gp_cov_scale(gp_cov_jitter(gp_cov_se(D2_ww)));
% $$$ theta_w{1} = [1; 4];
% $$$ covfunc_w{2} = gp_cov_scale(gp_cov_jitter(gp_cov_se(D2_ww)));
% $$$ theta_w{2} = [1; 3];
% $$$ covfunc_w{3} = gp_cov_scale(gp_cov_jitter(gp_cov_pp(sqrt(D2_ww),1)));
% $$$ theta_w{3} = [1; 2];
% $$$
% $$$ % GP module with component-wise factorization
% $$$ W_module = factor_module_gp_factorized(N_w, covfunc_w, theta_w, ...
% $$$                                        'update_hyperparameters', 2);
% $$$
% $$$ %
% $$$ % Isotropic noise module
% $$$ %
% $$$
% $$$ %noise_module = noise_module_fixed(1/s^2 * ones(N_w, N_x));
% $$$ noise_module = noise_module_isotropic(N_w, N_x, 1e-3, 1e-3, 'init', 100);
% $$$
% $$$ %
% $$$ % VB inference
% $$$ %
% $$$
% $$$ %
% $$$ % TODO:
% $$$ %
% $$$ % - pseudo inputs
% $$$ %
% $$$ % - put the noise to Q(W)
% $$$ %
% $$$ % - test more components
% $$$ %
% $$$ % - rotation, learn the hyperparameters jointly?
% $$$ %
% $$$ % - weighted noise
% $$$ %
% $$$
% $$$ Q = gpfa(D, Y, W_module, X_module, noise_module, ...
% $$$          'maxiter', 30, ...
% $$$          'update_noise', 2, ...
% $$$          'rotate', false);
% $$$
% $$$ noise_std = Q.Tau(1)^(-0.5)
% $$$
% $$$ tsplot(Q.X)
% $$$ tsplot(Q.W)
% $$$ %tsplot(Y, 'k+')
% $$$
% $$$ figure
% $$$ plot(Q.loglikelihood)
% $$$
% $$$ recon_error = rmse(Y_noiseless, Q.W'*Q.X)


function [err, errK] = getL2error3CR(node,elem,uvecexact,uh,quadOrder)
%% GETL2ERROR3CR L2 norm of the approximation error for vectorial CR element
%    need Matlab ver > 2016b for the vectorization to work
%
%    Reference: Error analysis of a decoupled finite element method for quad-curl problems
%               https://arxiv.org/abs/2102.03396
%
%    see alos: getL2error3
%
% Copyright (C) Long Chen. See COPYRIGHT.txt for details.

if ~exist('quadOrder','var'); quadOrder = 3; end
elem2face = dof3face(elem);

%%
NT = size(elem,1);
err = zeros(NT,1);
[lambda,weight] = quadpts3(quadOrder);
phi = 1-3*lambda;

nQuad = size(lambda,1);
for p = 1:nQuad
    uhp = uh(elem2face(:,1),:)*phi(p,1) + ...
        uh(elem2face(:,2),:)*phi(p,2) + ...
        uh(elem2face(:,3),:)*phi(p,3) + ...
        uh(elem2face(:,4),:)*phi(p,4);
    % quadrature points in the x-y coordinate
    pxyz = lambda(p,1)*node(elem(:,1),:) ...
         + lambda(p,2)*node(elem(:,2),:) ...
         + lambda(p,3)*node(elem(:,3),:) ...
         + lambda(p,4)*node(elem(:,4),:);
    err = err + weight(p)*sum((uvecexact(pxyz) - uhp).^2,2);
end

%%
d12 = node(elem(:,2),:)-node(elem(:,1),:);
d13 = node(elem(:,3),:)-node(elem(:,1),:);
d14 = node(elem(:,4),:)-node(elem(:,1),:);
volume = abs(dot(mycross(d12,d13,2),d14,2)/6);
err(isnan(err)) = 0; % singular point is excluded
err = volume.*err;
if nargout > 1; errK = sqrt(err); end
err = sqrt(sum(err));

end


function output = MSE(predicted, actual)
 output = mean(sum(((predictions-actual).^2)'));
end


function [u,edge,eqn,info] = Maxwellsaddle(node,elem,HB,pde,bdFlag,option)
%% MAXWELL Maxwell equation: lowest order edge element.
%
% u = Maxwellsaddle(node,elem,HB,pde,bdFlag) produces the lowest order edge
%   element approximation of the electric field of the time harmonic
%   Maxwell equation.
%
%                     curl(mu^(-1)curl u)  = J    in \Omega,
%                                  div  u  = 0
%                                   n � u = n � g_D  on \Gamma_D,
%                     n � (mu^(-1)curl u) = n � g_N  on \Gamma_N.
%
% based on the weak formulation
%
% (mu^{-1}curl u, curl v) +(grad p ,v)= (J,v) - <n � g_N,v>_{\Gamma_N}.
% (u,grad q)                          = 0
% Assume P\in H_0^1.
% The data of the equation is enclosed in the pde structure:
%   - pde.mu      : permeability, i.e., magnetic constant/tensor
%   - pde.omega   : wave number
%   - pde.J       : current density
%   - pde.g_D     : Dirichlet boundary condition
%   - pde.g_N     : Neumann boundary condition
%
% The mesh is given by (node,elem) and HB is needed for fast solvers. The
% boundary faces is specified by bdFlag; see bddoc.
%
% The function Maxwell assembes the matrix equation (A-M)*u = b and solves
% it by the direct solver (small size dof <= 2e3) or the HX preconditioned
% Krylov iterative methods (large size dof > 2e3).
%
% u = Maxwellsaddle(node,elem,HB,pde,bdFlag,option) specifies the solver options.
%   - option.solver == 'direct': the built in direct solver \ (mldivide)
%   - option.solver == 'mg':     multigrid-type solvers mg is used.
%   - option.solver == 'notsolve': the solution u = u_D.
% The default setting is to use the direct solver for small size problems
% and multigrid solvers for large size problems. For more options on the
% multigrid solver mg, type help mg.
%
% [u,edge] = Maxwellsaddle(node,elem,pde,bdEdge) returns also the edge array
% which is essential for edge elements.
%
% [u,edge,eqn] = Maxwellsaddle(node,elem,pde,bdEdge) returns also the equation
% structure eqn, which includes:
% - eqn.A: matrix for differential operator;
% - eqn.M: mass matrix;
% - eqn.f: right hand side
% - eqn.g: vector enclosed the Neumann boundary condition
%
% [u,edge,eqn,info] = Maxwellsaddle(node,elem,pde,bdEdge) returns also the
% information on the assembeling and solver, which includes:
% - info.assembleTime: time to assemble the matrix equation
% - info.solverTime:   time to solve the matrix equation
% - info.itStep:       number of iteration steps for the mg solver
% - info.error:        l2 norm of the residual b - A*u
% - info.flag:         flag for the mg solver.
%   flag = 0: converge within max iteration
%   flag = 1: iterated maxIt times but did not converge
%   flag = 2: direct solver
%   flag = 3: no solve
%
% Example
%   cubeMaxwell
%
% See also Maxwell1, Maxwell2, cubeMaxwell, mgMaxwell
%
% Reference page in Help browser
%       Maxwelldoc
%
% Created by Jie Zhou based on Maxwell(node,elem,pde,bdEdge) on
% 09,Sep,2013.
% Copyright (C) Long Chen. See COPYRIGHT.txt for details.

%% Set up optional input arguments
if ~exist('bdFlag','var'), bdFlag = []; end
if ~exist('option','var'), option = []; end

%% Sort elem to ascend ordering
[elem,bdFlag] = sortelem3(elem,bdFlag);

%% Construct Data Structure
[elem2dof,edge] = dof3edge(elem);
locEdge = [1 2; 1 3; 1 4; 2 3; 2 4; 3 4];
N = size(node,1);   NT = size(elem,1);  Ne = size(edge,1);

%% Compute coefficients
if ~isfield(pde,'mu'), pde.mu = 1; end
if ~isempty(pde.mu) && isnumeric(pde.mu)
    mu = pde.mu;                % mu is an array
else                            % mu is a function
    center = (node(elem(:,1),:) + node(elem(:,2),:) + ...
              node(elem(:,3),:) + node(elem(:,4),:))/4;
    mu = pde.mu(center);
end
if ~isfield(pde,'epsilon'), pde.epsilon = 0; end
if ~isempty(pde.epsilon) && isnumeric(pde.epsilon)
    epsilon = pde.epsilon;      % epsilon is an array
else                            % epsilon is a function
    center = (node(elem(:,1),:) + node(elem(:,2),:) + ...
              node(elem(:,3),:) + node(elem(:,4),:))/4;
    epsilon = pde.epsilon(center);
end
if isfield(pde,'omega')
    omega = pde.omega;
else
    omega = 1;
end
epsilon = omega^2*epsilon;

tstart = tic;
%% Element-wise basis
% edge indices of 6 local bases:
% [1 2], [1 3], [1 4], [2 3], [2 4], [3 4]
% phi = lambda_iDlambda_j - lambda_jDlambda_i;
% curl phi = 2*Dlambda_i � Dlambda_j;
[Dlambda,volume] = gradbasis3(node,elem);
curlPhi(:,:,6) = 2*mycross(Dlambda(:,:,3),Dlambda(:,:,4),2);
curlPhi(:,:,1) = 2*mycross(Dlambda(:,:,1),Dlambda(:,:,2),2);
curlPhi(:,:,2) = 2*mycross(Dlambda(:,:,1),Dlambda(:,:,3),2);
curlPhi(:,:,3) = 2*mycross(Dlambda(:,:,1),Dlambda(:,:,4),2);
curlPhi(:,:,4) = 2*mycross(Dlambda(:,:,2),Dlambda(:,:,3),2);
curlPhi(:,:,5) = 2*mycross(Dlambda(:,:,2),Dlambda(:,:,4),2);
DiDj = zeros(NT,4,4);
for i = 1:4
    for j = i:4
        DiDj(:,i,j) = dot(Dlambda(:,:,i),Dlambda(:,:,j),2);
        DiDj(:,j,i) = DiDj(:,i,j);
    end
end

%% Assemble matrices
ii = zeros(21*NT,1); jj = zeros(21*NT,1);
sA = zeros(21*NT,1); sM = zeros(21*NT,1);
index = 0;
for i = 1:6
    for j = i:6
        % local to global index map
        % curl-curl matrix
        Aij = dot(curlPhi(:,:,i),curlPhi(:,:,j),2).*volume./mu;
        ii(index+1:index+NT) = double(elem2dof(:,i));
        jj(index+1:index+NT) = double(elem2dof(:,j));
        sA(index+1:index+NT) = Aij;
        % mass matrix
        % locEdge = [1 2; 1 3; 1 4; 2 3; 2 4; 3 4];
        i1 = locEdge(i,1); i2 = locEdge(i,2);
        j1 = locEdge(j,1); j2 = locEdge(j,2);
        Mij = 1/20*volume.*( (1+(i1==j1))*DiDj(:,i2,j2) ...
                           - (1+(i1==j2))*DiDj(:,i2,j1) ...
                           - (1+(i2==j1))*DiDj(:,i1,j2) ...
                           + (1+(i2==j2))*DiDj(:,i1,j1));
        %Mij = Mij.*epsilon;
        sM(index+1:index+NT) = Mij;
        index = index + NT;
    end
end
clear curlPhi % clear large size data
diagIdx = (ii == jj);   upperIdx = ~diagIdx;
A = sparse(ii(diagIdx),jj(diagIdx),sA(diagIdx),Ne,Ne);
AU = sparse(ii(upperIdx),jj(upperIdx),sA(upperIdx),Ne,Ne);
A = A + AU + AU';
M = sparse(ii(diagIdx),jj(diagIdx),sM(diagIdx),Ne,Ne);
MU = sparse(ii(upperIdx),jj(upperIdx),sM(upperIdx),Ne,Ne);
M = M + MU + MU';
grad =icdmat(double(edge),[-1,1]);
G    = grad'*M;

% bigA = A + M;

%% Assemble right hand side
f = zeros(Ne,1);
if ~isfield(pde,'J') || (isfield(pde,'J') && isreal(pde.J) && all(pde.J==0))
    pde.J = [];
end
if ~isfield(option,'fquadorder')
    option.fquadorder = 2;   % default order is 3
end
if isfield(pde,'J') && ~isempty(pde.J)
    [lambda,w] = quadpts3(option.fquadorder);
    nQuad = size(lambda,1);
    bt = zeros(NT,6);
    for p = 1:nQuad
        % quadrature points in the x-y-z coordinate
        pxyz = lambda(p,1)*node(elem(:,1),:) ...
             + lambda(p,2)*node(elem(:,2),:) ...
             + lambda(p,3)*node(elem(:,3),:) ...
             + lambda(p,4)*node(elem(:,4),:);
        Jp = pde.J(pxyz);
    %   locEdge = [1 2; 1 3; 1 4; 2 3; 2 4; 3 4];
        for k = 1:6
            i = locEdge(k,1); j = locEdge(k,2);
            % phi_k = lambda_iDlambda_j - lambda_jDlambda_i;
            phi_k = lambda(p,i)*Dlambda(:,:,j)-lambda(p,j)*Dlambda(:,:,i);
            rhs = dot(phi_k,Jp,2);
            bt(:,k) = bt(:,k) + w(p)*rhs;
        end
    end
    bt = bt.*repmat(volume,1,6);
    f = accumarray(elem2dof(:),bt(:),[Ne 1]);
end
clear pxyz Jp bt rhs phi_k

%% Set up solver
if isempty(option) || ~isfield(option,'solver')    % no option.solver
    if Ne <= 1e4  % Direct solver for small size systems
        option.solver = 'direct';
    else            % Multigrid-type  solver for large size systems
        option.solver = 'cg';
    end
end
solver = option.solver;

%% Assembeling corresponding matrices for HX preconditioner
if ~strcmp(solver,'direct') && ~strcmp(solver,'nosolve')
    AP = sparse(N,N);  % AP = - div(mu^{-1}grad) + |Re(epsilon)| I
    BP = sparse(N,N);  % BP = - div(|Re(epsilon)|grad)
    for i = 1:4
        for j = i:4
            temp = DiDj(:,i,j).*volume;
            Aij = 1./mu.*temp;
            Bij = abs(real(epsilon)).*temp;
            Mij = 1/20*abs(real(epsilon)).*volume;
            if (j==i)
                AP = AP + sparse(elem(:,i),elem(:,j),Aij+2*Mij,N,N);
                BP = BP + sparse(elem(:,i),elem(:,j),Bij,N,N);
            else
                AP = AP + sparse([elem(:,i);elem(:,j)],[elem(:,j);elem(:,i)],...
                                 [Aij+Mij; Aij+Mij],N,N);
                BP = BP + sparse([elem(:,i);elem(:,j)],[elem(:,j);elem(:,i)],...
                                 [Bij; Bij],N,N);
            end
        end
    end
end
clear Aij Bij Mij

%% Boundary conditions
if ~isfield(pde,'g_D'), pde.g_D = []; end
if ~isfield(pde,'g_N'), pde.g_N = []; end
if ~isfield(pde,'g_R'), pde.g_R = []; end
if (isempty(pde.g_D) && isempty(pde.g_N) && isempty(pde.g_R))
    % no boundary data is given = homogenous Neumann boundary condition
    bdFlag = [];
end

%% Part 1: Find Dirichlet dof and modify the matrix
% Find Dirichlet boundary dof: fixedDof
isBdEdge = [];
if isempty(bdFlag) && ~isempty(pde.g_D) && isempty(pde.g_N)
    % Dirichlet boundary condition only
    bdFlag = setboundary3(node,elem,'Dirichlet');
end
if ~isempty(bdFlag)
    % Find boundary edges and nodes
    isBdEdge = false(Ne,1);
    isBdNode = false(N,1);
    isBdEdge(elem2dof(bdFlag(:,1) == 1,[4,5,6])) = true;
    isBdEdge(elem2dof(bdFlag(:,2) == 1,[2,3,6])) = true;
    isBdEdge(elem2dof(bdFlag(:,3) == 1,[1,3,5])) = true;
    isBdEdge(elem2dof(bdFlag(:,4) == 1,[1,2,4])) = true;
    bdEdge = edge(isBdEdge,:);
    isBdNode(bdEdge) = true;
end
% modify the matrix to include the Dirichlet boundary condition
if any(isBdEdge)  % contains Dirichlet boundary condition
    bdidx = zeros(Ne,1);
    bdidx(isBdEdge) = 1;
    Tbd = spdiags(bdidx,0,Ne,Ne);
    Te = spdiags(1-bdidx,0,Ne,Ne);
    AD = Te*(A - epsilon*M)*Te + Tbd;
    if ~strcmp(solver,'direct') && ~strcmp(solver,'nosolve')
        % modify the corresponding Poisson matrix
        bdidx = zeros(N,1);
        bdidx(isBdNode) = 1;
        Tbd = spdiags(bdidx,0,N,N);
        Tv = spdiags(1-bdidx,0,N,N);
        AP = Tv*AP*Tv + Tbd;
        BP = Tv*BP*Tv + Tbd;
       % G  = Te*G*Tv  ;
    end
else      % pure Neumann boundary condition
    AD = A - epsilon*M  ;
    if ~strcmp(solver,'direct') && ~strcmp(solver,'nosolve')
       BP = BP + 1e-8*speye(N);  % make B non-singular
    end
end

%% Part 2: Find boundary edges and modify the load b
g = zeros(Ne,1);
% Find Neumann boundary faces
if isempty(bdFlag) && (~isempty(pde.g_N) || ~isempty(pde.g_R))
    bdFlag = setboundary3(node,elem,'Neumann');
end
% non-zero Neumann boundary condition
if ~isempty(bdFlag) && ~isempty(pde.g_N)
    % face 1
    isBdElem = find(bdFlag(:,1) == 2); %#ok<*NASGU>
    face = [2 3 4]; face2locdof = [6 5 4];
    if ~isempty(isBdElem)
        bdb = bdfaceintegral(isBdElem,face,face2locdof);
        g = bdb;
    end
    % face 2
    isBdElem = find(bdFlag(:,2) == 2);
    face = [1 4 3]; face2locdof = [6 2 3];
    if ~isempty(isBdElem)
        bdb = bdfaceintegral(isBdElem,face,face2locdof);
        g = g + bdb;
    end
    % face 3
    isBdElem = find(bdFlag(:,3) == 2);
    face = [1 2 4]; face2locdof = [5 3 1];
    if ~isempty(isBdElem)
        bdb = bdfaceintegral(isBdElem,face,face2locdof);
        g = g + bdb;
    end
    % face 4
    isBdElem = find(bdFlag(:,4) == 2);
    face = [1 3 2]; face2locdof = [4 1 2];
    if ~isempty(isBdElem)
        bdb = bdfaceintegral(isBdElem,face,face2locdof);
        g = g + bdb;
    end
    f = f - g;
end
% nonzero Dirichlet boundary condition
u = zeros(Ne,1);
if ~isempty(bdEdge) && ~isempty(pde.g_D) && ...
   ~(isnumeric(pde.g_D) && all(pde.g_D == 0))
    % else no bddof or g_D = 0 (no modification needed)
    if (isnumeric(pde.g_D) && length(pde.g_D) == Ne)
        u(isBdEdge) = pde.g_D(isBdEdge);
    else
        u(isBdEdge) = edgeinterpolate(pde.g_D,node,bdEdge);
    end
    f = f - (A - epsilon*M)*u;
    f(isBdEdge) = u(isBdEdge);
end
%% We always assume the Langrange multer is in zeros boundary condition
p     = zeros(N,1);
g0    = zeros(N,1);
g0    = -G*u;

%% Remark
% The order of assign Neumann and Dirichlet boundary condition is
% important to get the right setting of the intersection of Dirichlet and
% Neumann faces.

%% Record assembling time
info.assembleTime = toc(tstart);
if ~isfield(option,'printlevel'), option.printlevel = 1; end
if option.printlevel >= 1
    fprintf('Time to assemble matrix equation %4.2g s\n',info.assembleTime);
end

%% Solve the system of linear equations
% if strcmp(solver,'direct')
%     % exact solver
%     tstart = tic;
%     freeDof = find(~isBdEdge);
%     u(freeDof) = AD(freeDof,freeDof)\f(freeDof);
%     time = toc(tstart); itStep = 0; flag = 2; err = norm(f - AD*u);
% elseif strcmp(solver,'nosolve')
%     eqn = struct('A',A,'M',M,'f',f,'g',g,'bigA',AD,'isBdEdge',isBdEdge);
%     info = [];
%     return;
% else
% %     u0 = edgeinterpolate(pde.g_D,node,edge);
%     u0 = u;
%     option.x0 = u0;
%     [u,flag,itStep,err,time] = mgMaxwell(AD,f,AP,BP,node,elem,edge,HB,isBdEdge,option);
% end
temp  = zeros(N+Ne,1);
bigAD = [AD G';G sparse(N,N)];
freeEdge = find(~isBdEdge);
freeNode = find(~isBdNode);
freeDof  = [freeEdge;Ne+freeNode];
tstart = tic;
temp(freeDof) = bigAD(freeDof,freeDof)\[f(freeEdge);g0(freeNode)];
time = toc(tstart);itStep = 0; flag = 2; err = norm([f(freeEdge);g0(freeNode)] - bigAD(freeDof,freeDof)*temp(freeDof));
u(freeEdge) = temp(freeEdge);
p(freeNode) = temp(Ne + freeNode);

%check the Langrange is correct or not.
normp  = norm(p);
if(normp>1.0/N)
disp('the Langrange multer is wrong')
end

%% Output
eqn = struct('A',A,'M',M,'f',f,'g',g,'bigA',AD,'isBdEdge',isBdEdge);
info = struct('solverTime',time,'itStep',itStep,'error',err,'flag',flag);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% subfunctions bdfaceintegral
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    function bdb = bdfaceintegral(isBdElem,face,face2locdof)
    %% Compute boundary surface integral of lowest order edge element.
    %  bdb(k) = \int_{face} (n�g_N, phi_k) dS

    %% Compute scaled normal
    faceIdx = true(4,1);
    faceIdx(face) = false;
    normal = -3*repmat(volume(isBdElem),1,3).*Dlambda(isBdElem,:,faceIdx);

    %% Data structure
    tetLocEdge = [1 2; 1 3; 1 4; 2 3; 2 4; 3 4]; % edge of a tetrahedral [1 2 3 4]
    face2locEdge = [2 3; 3 1; 1 2]; % edge of the face [1 2 3]

    %% Compute surface integral
    Nbd = length(isBdElem);
    bt = zeros(Nbd,3);
    idx = zeros(Nbd,3,'int32');
    [lambda,w] = quadpts(3); % quadrature order is 3
    nQuad = size(lambda,1);
    for pp = 1:nQuad
        % quadrature points in the x-y-z coordinate
        pxyz = lambda(pp,1)*node(elem(isBdElem,face(1)),:) ...
             + lambda(pp,2)*node(elem(isBdElem,face(2)),:) ...
             + lambda(pp,3)*node(elem(isBdElem,face(3)),:);
        gNp = pde.g_N(pxyz,normal);
        for s = 1:3
            kk = face2locdof(s);
            pidx = face(face2locEdge(s,1))< face(face2locEdge(s,2));
            % phi_k = lambda_iDlambda_j - lambda_jDlambda_i;
            % lambda_i is associated to the local index of the face [1 2 3]
            % Dlambda_j is associtated to the index of tetrahedron
            % - when the direction of the local edge s is consistent with the
            %   global oritentation given in the triangulation,
            %           s(1) -- k(1),  s(2) -- k(2)
            % - otherwise
            %           s(2) -- k(1),  s(1) -- k(2)
            if pidx
                phi_k = lambda(pp,face2locEdge(s,1))*Dlambda(isBdElem,:,tetLocEdge(kk,2)) ...
                      - lambda(pp,face2locEdge(s,2))*Dlambda(isBdElem,:,tetLocEdge(kk,1));
            else
                phi_k = lambda(pp,face2locEdge(s,2))*Dlambda(isBdElem,:,tetLocEdge(kk,2)) ...
                      - lambda(pp,face2locEdge(s,1))*Dlambda(isBdElem,:,tetLocEdge(kk,1));
            end
            rhs = dot(phi_k,gNp,2);
            bt(:,s) = bt(:,s) + w(pp)*rhs; % area is included in normal; see line 28
            idx(:,s) = elem2dof(isBdElem,kk);
        end
    end
    %% Distribute to DOF
    bdb = accumarray(idx(:),bt(:),[Ne 1]);
    end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
end