% FID from a 1D-periodic system including RF
% the use of symmetry can be disabled by setting testmode to a non-zero value

M=2; % spins per cell
N=3; % number of cells

r=1.5; % internal distance (A)
R=3; % distance between cells

testmode=0; % set to do calculation with/without using symmetry

vRF=0; % amplitude of RF (phase is zero);

inccs=0; % include shifts?
shifts=[0 0]; % (isotropic) chemical shifts of spins

thet=45*pi/180.0; % angle (relative to z)

maxorder=floor(N/2); % maximum coupling 'order'

% calculate positions of spins (in xz plane)
baseco=zeros(M,3);
baseco(1,:)=[ 0 0 0];
if M>1
  baseco(2,:)=[ r*cos(thet) 0 r*sin(thet)];
  if M>2
     baseco(3,:)=[r*cos(-thet) 0 r*sin(-thet)];
  end
end

sw=20000; % spectral width for histogram
bins=512; % number of bins in histogram
Hist=zeros(1,bins);

disp('Co-ordinates'); % show nuclei positions
disp(baseco);

vector=[ R 0 0 ]; % intercell vector

total=M*N; % total number of spins

distances=zeros(M,total);
couplings=zeros(M,total);

incrf=(vRF~=0); % whether to include RF or not

for j=1:M
  for k=1:M % loop over spin pairs
     iv=baseco(j,:)-baseco(k,:); % internuclear vector (within cell)
     for celln=0:N-1 % loop over cells
        if ( (celln~=0) | (j~=k) ) % don't include coupling of spin with itself!
           tb=iv+celln*vector;
           if (2*celln>N) % choose shortest distance (in cells) between counting cells forward, and cells backward
              tb=iv-(N-celln)*vector;
           else
              if ( (2*celln==N) & (k<j)) % ambiguous for N even, choose specific spin order (k<j)
                 tb=-iv+celln*vector;
              end
           end
     
           [th,phi,r]=cart2sph(tb(1),tb(2),tb(3)); % convert to spherical co-ords
           
           sk=celln*M+k; % convert from cell,spin pair to spin index
           distances(j,sk)=r; % distance between spins
           
           if ( (celln<=maxorder) | (N-celln<=maxorder) ) % use coupling if within max order

              %d=rtod('13C','13C',r);
              d=7.5936e3/(r^3); % dipolar coupling (based on 13C-13C distances in A);
              % would introduce rotation from crystallite powder angle here
              couplings(j,sk)=d;
           end
        end
     end
  end
end

disp('Internuclear distances');
disp(distances);

disp('Couplings');
disp(couplings);

dim=2^total; % dimension of Hilbert space
sys=repmat(0.5,1,total); % spin system

if testmode~=0
   % calculate dipolar coupling Hamiltonian
   
   Hsys=0; % couplings within cells
   for j=1:M-1
      for k=j+1:M
         for celln=0:N-1
            sj=celln*M+j;
            sk=celln*M+k;
            Hsys=Hsys+couplings(j,k)*(2*S(sys,sj,'z')*S(sys,sk,'z')-S(sys,sj,'x')*S(sys,sk,'x')-S(sys,sj,'y')*S(sys,sk,'y'));
         end
      end
   end
   
   HsysExt=0; % couplings between cells
   for celln=1:N-1
      for j=1:M
         for k=1:M
            d=0.5*couplings(j,celln*M+k);
            for cell2=0:N-1
               sk=cell2*M+k;
               cellf=mod(cell2+celln,N);
               sj=cellf*M+j;
               HsysExt=HsysExt+d*(2*S(sys,sj,'z')*S(sys,sk,'z')-S(sys,sj,'x')*S(sys,sk,'x')-S(sys,sj,'y')*S(sys,sk,'y'));
            end
         end
      end
   end

   Hcs=0; % chemical shift Hamiltonian
   if inccs~=0
      for j=1:M
         Hcs=Hcs+shifts(j)*F(sys,[j:M:total],'z');
      end   
   end
      
   Hrf=vRF*F(sys,'x'); % RF Hamiltonian
   
   Htotal=Hsys+HsysExt+Hcs+Hrf; % total Hamiltonian
   
   %disp('Hamiltonian');
   %disp(Htotal);
   
   sigma0=F(sys,'+'); % initial density matrix
   
   %disp('sigma0');
   %disp(sigma0);
   
   starttime=cputime;
   Hist=zeros(1,bins);
   Hist=calcsignal(Hist,sw/bins,Htotal,sigma0); % calculate spectral histogram
   
else % use symmetry
   
   useeigs=N; % number of eigenvalues used
   
   startk=1;
   endk=useeigs; % change these to focus on particular k values
   
   DEADBEEF=1e30; % obviously invalid value (handy for debugging!)

   % set up arrays that depend only on N
   eigfacs=zeros(1,N-1); % pre-calculate exp(i*k) factors
   for j=1:N-1
      eigfacs(j)=exp(2*pi*i*j/N);
   end
   
   scalefacs=repmat(DEADBEEF,N,N); % N/sqrt(n_A n_B) block scaling factors
   haseig=zeros(N,useeigs); % Does block with given size contain a given k value?
   
   for j=1:N
      vec=[1:N];
      if mod(N,j)==0
         step=N/j;
         haseig(j,step:step:N)=1;
         scalefacs(j,:)=N./sqrt(j*vec);
      end
   end

   % find sets of symmetry-linked states
   used=zeros(1,dim);
   count=dim; % states used
   mask=2^M-1;
   T={};
   while count~=0
      start=1; % find next unused state
      while used(start)
         start=start+1;
      end
      st_state=start-1; % subtract 1 to convert into state
      
      cstate=st_state;
      state_list=[]; % set of translation-linked states
      while 1
         used(cstate+1)=1;
         state_list=[ state_list cstate];
         cstate=bitor(bitshift(cstate,-M),bitshift(bitand(cstate,mask),total-M));
         if cstate==st_state
            break;
         end
      end
      T={ T{:} state_list};
      count=count-length(state_list); % decrease by no. of states used
   end
   %disp('Set structure');
   %for k=1:length(T)
   %   disp(T{k});
   %end
     
   totblks=length(T); % number of sets
   
   % 'reverse' indices so that when we "generate" a new state
   % by bit flipping, we know where it belongs
   state_to_block=zeros(1,dim); % converts index to set number
   state_to_index=zeros(1,dim); % converts index to position within set
   
   mz2s=zeros(1,totblks); % 2*mz value for each set
   blksizes=zeros(1,totblks); % size of each set
   
   for blk=1:totblks
      clist=T{blk};
      blksizes(blk)=length(clist);
      mz2s(blk)=floor(0.5+2*Fzel(clist(1),total)); % 2*mz value
      for k=1:length(clist)
         state_to_block(clist(k)+1)=blk;
         state_to_index(clist(k)+1)=k;
      end
   end
   
   eigcount=zeros(1,N); % size of each k block
   eigptrs=zeros(totblks,N); % index into each block of k value for given set (zero if absent)
   % note that k=N is used for k=0 (Matlab indices starting at 1)
   
   for j=1:totblks
      blk=T{j};
      size=length(blk);
      step=N/size;
      for k=step:step:N
         eigcount(k)=eigcount(k)+1; % inc eigenvalue count
         eigptrs(j,k)=eigcount(k);
      end
   end
   
   rfplus=vRF/2.0;
   rfminus=vRF/2.0; % RF terms in +/- blocks of Hamiltonian

   % find symmetrised dipolar Hamiltonian
   
   % create empty array of matrices for symmetrised Hamiltonian and density matrix
   Hsym=cell(1,useeigs);
   sigmasym=cell(1,useeigs);
   for j=1:useeigs
      dj=eigcount(j);
      Hsym{j}=zeros(dj,dj);
      sigmasym{j}=Hsym{j};
   end
   
   for j=1:totblks % loop ovr 'row' sets
      reigptrs=eigptrs(j,:);
      rblksize=blksizes(j); % no. of bra states
      %rstep=N/rblksize;
      for k=1:totblks
         coher=(mz2s(j)-mz2s(k))/2; % coherence order
         
         if abs(coher)<2 % only coherences needed are 0 (system Hamiltonian) and 1 (RF Hamiltonian and density matrix)
            cblksize=blksizes(k); % number of ket states
            %cstep=N/cblksize;
            ceigptrs=eigptrs(k,:);
            %step=max(cstep,rstep);
            H=repmat(0,rblksize,cblksize); % zero matrices
            if coher==1 % only need +1 coherence order
               sigma=H;
            end
            cstates=T{k}; % ket states
            
            switch coher
            case 0 % ZQ block
               for spin1=1:M % loop over spins within cell
                  
                  maskj=bitshift(1,total-spin1); % mask to extract state of spin1
                  
                  % add chemical shift
                  if (inccs~=0) & (j==k) % diagonal blocks only
                     for jj=1:cblksize
                        if bitand(cstates(jj),maskj)~=0 % get state of spin1
                           izval=-0.5; % 1 is beta (mz=-1/2) state
                        else
                           izval=0.5; % if alpaha state
                        end
                        H(jj,jj)=H(jj,jj)+izval*shifts(spin1);
                     end
                  end
                  
                  for spin2=1:M % loop over spin pairs
                     for celln=0:N-1 % loop over unit cells (0 same cell)
                        if (spin1>spin2) | (celln~=0) % skip spin1<=spin2 within cell
                           sk=M*celln+spin2; % convert to spin index
                           masksk=bitshift(1,total-sk); % mask for state of 'other' spin
                           if celln 
                              coup=0.25*couplings(spin1,sk); % scaling factor of 0.5 for couplings outside cell (otherwise couplings are double counted)
                           else
                              coup=0.5*couplings(spin1,sk);
                           end
                           mask=bitor(masksk,maskj); % flip mask for spin pair
                           for jj=1:cblksize % loop over ket states
                              cstate=cstates(jj);
                              if j==k % zz terms (only in diagonal blocks)
                                 if xor(bitand(cstate,maskj),bitand(cstate,masksk))
                                    H(jj,jj)=H(jj,jj)-coup; % one spin up, one spin down
                                 else
                                    H(jj,jj)=H(jj,jj)+coup; % both spins in same state
                                 end
                              end
                              nstate=bitxor(cstate,mask); % new state from flipping both spins
                              if state_to_block(nstate+1)==j % is new state in set j ?
                                 destr=state_to_index(nstate+1); % get index of (symmetrised) state
                                 H(destr,jj)=H(destr,jj)-coup; % xx, yy terms
                              end
                           end
                        end
                     end
                  end
               end   
            case {1,-1} % +/-1 terms
               for jj=1:cblksize % loop over ket states
                  cstate=cstates(jj);
                  if coher==1 % +1 coherence
                     for spin=1:M
                        flipmask=bitcmp(bitshift(1,M-spin),total);
                        nstate=bitand(cstate,flipmask); % clear bit (set to alpha state)
                        if state_to_block(nstate+1)==j % is new state in set j? (note that if spin was already in alpha state, "new" state cannot be in set j)
                           where=state_to_index(nstate+1); % index of (symmetrised) state
                           H(where,jj)=rfplus;
                           sigma(where,jj)=1.0; % F-
                        end
                     end
                  else
                     for spin=1:M
                        flipmask=bitshift(1,M-spin);
                        nstate=bitor(cstate,flipmask); % set bit (set to beta state)
                        if state_to_block(nstate+1)==j % is new state in set j?
                           H(state_to_index(nstate+1),jj)=rfminus;
                        end
                     end
                  end
               end
            end
               
            scalef=scalefacs(rblksize,cblksize); % sqrt(N/(rblksize*cblksize));
            
            for m=startk:endk % loop over k values of interest
               if (haseig(cblksize,m)~=0) & (haseig(rblksize,m)~=0) % is there a symmetrised state with this value of k?
                  s=0.0;
                  sig=0.0;
                  for ii=1:rblksize % loop over (symmetrised) bra states
                     for jj=1:cblksize % loop over (symmetrised) ket states
                        wh=mod(m*(jj-ii),N); % exp(i*k*(col-row)/N)
                        if wh~=0
                           s=s+H(ii,jj)*eigfacs(wh); % accumulate Hamiltonian element
                           if coher==1
                              sig=sig+sigma(ii,jj)*eigfacs(wh); % accumulate density matrix element
                           end
                        else
                           s=s+H(ii,jj);
                           if coher==1
                              sig=sig+sigma(ii,jj);
                           end
                        end
                     end
                  end
                  % store element in correct place in eigenvalue blocks
                  Hsym{m}(reigptrs(m),ceigptrs(m))=s*scalef;
                  if coher==1
                     sigmasym{m}(reigptrs(m),ceigptrs(m))=sig*scalef;
                  end
               end
            end
         end
      end
   end
   
   starttime=cputime;
   Hist=zeros(1,bins);
   for k=startk:endk % loop over k values of interest
      disp(sprintf('k=%i',k));
      
      HH=(Hsym{k}+Hsym{k}')/2; % force Hamiltonian to be Hermitian
      Hist=calcsignal(Hist,sw/bins,HH,sigmasym{k}); % accumulate spectral histogram
   end
end  

%disp(sprintf('Time taken for signal calculation: %g',cputime-starttime));

fscale=[-bins/2+1:bins/2]*(sw/1000)/bins;
bar(fscale,Hist);
xlabel('frequency / kHz');

if testmode~=0
   title('Without symmetry');
else
   title('With symmetry');
end