classdef mga_golden_shots
  % Class for creating/handling set of mga golden shots
  % These shots serve as reference for desired mga behaviour
  % used in tests by test_mga_golden_shots
  
  properties
    table_filename = 'golden_shot_table.csv';
    max_shots = 10; % maximum number shots
    source_shotlist % list of source shots
    first_shot = 6794; % first shot for which mgp(shot) works
    user_shot_offset = 100; % offset of user shots for writing
    ask_confirmation = true; % ask confirmation before overwriting user shots
    pp cell % cell array with parameters of interest or function handles
  end
  
  methods
    function obj = mga_golden_shots()
      % constructor
      [current_shot,s] = mdsopen(0);
      assert(isodd(s),'error opening shot');
      obj.source_shotlist = (current_shot-1):-1:obj.first_shot;
      obj.table_filename = fullfile(fileparts(mfilename('fullpath')),obj.table_filename);
      
      zp_low  = @(p)any(p.zmajo1 <-0.11);
      zp_high = @(p)any(p.zmajo1 > 0.11);
      zp_mid  = @(p)any(p.zmajo1 < 0.11 & p.zmajo1>-0.11);
      
      % parameters for which unique values are checked
      obj.pp = {'iscramb','inova','iohfb','if36fb','ierat','nfast','ikriz','mvloop','nvvel','amagic','nelz'};
      % function handles for which unique return results are checked
      obj.pp = [obj.pp,{zp_low,zp_mid,zp_high}]; % add function handles
      
    end
    
    function T = find_and_write_golden_shots(obj)
      % find shots with unique mgp parameters and write them to a shot range in the user's PCS tree
      %   T: Table of unique shots and their mgp parameter values
      
      T = obj.find_unique_mgp;
      T = obj.populate_user_shots(T);
    end
    
    function T = find_unique_mgp(obj)
      % find shots with unique mgp parameter combinations of interests
      
      shotlist = obj.source_shotlist;
      
      % Pre-allocate tables
      varnames = cell(1,numel(obj.pp));
      for ip = 1:numel(varnames)
        myp = obj.pp{ip};
        switch class(myp)
          case {'char','string'}
            varnames{ip} = myp;
          case 'function_handle'
            varnames{ip} = matlab.lang.makeValidName(func2str(myp)); % make function handle string into valid variable name
          otherwise
            error('can''t treat case %s',class(myp));
        end
      end
      varnames = ['golden_shot','silver_shot','source_shot',varnames]; % prepend golden and source shot number
      
      ncols = numel(varnames);
      T = array2table(zeros(obj.max_shots,ncols,'int32'),'VariableNames',varnames);
      
      itable = 0; % table row counter
      % loop over shots
      for ishot = 1:numel(shotlist)
        try
          shot = int32(shotlist(ishot));
          P = mgp(shot);
          
          % fill table row values
          tablevalues = zeros(1,numel(obj.pp)); % init
          for ip = 1:numel(obj.pp) % scan over parameters/function handles
            myp = obj.pp{ip};
            if isa(myp,'function_handle')
              tablevalues(ip) = myp(P); % call function handle on parameter structure
            else
              tablevalues(ip) = P.(myp); % value from field directly
            end
          end
          % new row to be added to table
          newrow = array2table([0,0,shot,tablevalues],'VariableNames',varnames);
          
          equalentries = ismember(newrow(:,4:end),T(:,4:end));
          if ~equalentries % check uniqueness
            itable = itable+1; % increment table index
            T(itable,:) = newrow; % add values to this table rowi
            
            fprintf('\nshot %d added as new unique parameter combination number %d\n',shot,itable)
            disp(T(itable,:))
          else
            isame = find(equalentries,1,'first');
            fprintf('shot %d has same parameters as %d\n',shot,T.source_shot(isame)),
            % already exists, skip
          end
          if itable>=obj.max_shots, break; end % exit when maximum shots reached
          
        catch ME
          fprintf('error parsing mgams parameters for shot: %d.\n Error message: %s\n',shot,getReport(ME));
        end
      end
      % remove any possibly zero table rows
      T = T(1:itable,:);
      
      % save to table csv file if not asked for T output
      if nargout==0
        writetable(T,obj.table_filename)
        fprintf('wrote table to %s\n',obj.table_filename);
      end
    end
    
    function T = populate_user_shots(obj,T,varargin)
      % populate shots in user space corresponding to list of shots in table
      % inputs:
      % T: table with golden shot data (if not passed, loaded from obj.table_filename
      % optional parameter-argument pairs:
      % golden_shots : list of specific golden shots from the table to rerun
      % silver_shots : list of specific silver shots from the table to rerun
      % source_shots : list of specific source shots from the table to rerun
      % table
      
      if nargin==1 || isempty(T) % read T from file if not passed explicitly
        T = readtable(obj.table_filename);
      end
      source_shots = T.source_shot;
      golden_shots = T.golden_shot;
      silver_shots = T.silver_shot;
      
      nsources = numel(source_shots);
      
      if nargin<=2
        i_populate = 1:nsources; % populate all
      else
        % choose which shots to rerun based on user input
        p=inputParser;
        p.addParameter('source_shots',[]);
        p.addParameter('golden_shots',[]);
        p.addParameter('silver_shots',[]);
        p.parse(varargin{:});
        P = p.Results;
        
        if ~isempty(P.source_shots)
          [check,i_source] = ismember(P.source_shots,source_shots);
          assert(all(check),'no source shot found in table matching user request')
        else
          i_source = [];
        end
        if ~isempty(P.golden_shots)
          [check,i_golden] = ismember(P.golden_shots,golden_shots);
          assert(all(check),'no golden shot found in table matching user request')
        else
          i_golden = [];
        end
        if ~isempty(P.silver_shots)
          [check,i_silver] = ismember(P.silver_shots,silver_shots);
          assert(all(check),'no silver shot found in table matching user request')
        else
          i_silver = [];
        end
        i_populate = unique([i_source(:);i_golden(:);i_silver(:)]).'; % Force row vector orientation
      end
      
      % get user allowed shot range
      myrange = mgu.myrange;
      shot_start = myrange{1}+obj.user_shot_offset;
      last_shot = myrange{2};
      
      assert(shot_start + nsources <= last_shot,...
        'number of shots to be written will exceed user writable shot range (%d-%d). Write fewer shots or choose lower shot offset',...
        shot_start,last_shot);
      
      % check before overwriting shot range
      if obj.ask_confirmation && (numel(i_populate) == nsources)
        prompt = sprintf('Warning, this will overwrite shots %d-%d, are you sure (y/n)',...
          shot_start,shot_start+2*nsources);
        x = input(prompt,'s');
        if ~contains(lower(x),'y'); fprintf('aborting\n'); return; end
      end
      
      % loop over source shots to be populated
      for fort_equiv = [false,true]
        for ishot = i_populate
          % source shot
          source_shot = source_shots(ishot);
          
          % destination user shot based on original table entry
          user_shot = shot_start + nsources*fort_equiv + ishot -1;
          % write shot
          s = obj.write(source_shot,user_shot,fort_equiv);
          
          if fort_equiv
            % MGA-FORTRAN comparison, silver shot
            T.silver_shot(ishot) = int32(user_shot*s);
          else
            % MGA-MATLAB comparison, golden shot
            T.golden_shot(ishot) = int32(user_shot*s);
          end
        end
      end
      
      % save updated table with golden shot numbers
      writetable(T,obj.table_filename);
      fprintf('wrote table to %s\n',obj.table_filename);
    end
  end
  
  methods(Static)
    % Static method to get list of shot numbers from table
    function shots = get_golden_shots(table_filename,nshots)
      % get vector of shot numbers where golden shots are stored
      T = readtable(table_filename);
      shots = T.golden_shot;
      shots = shots(shots~=0);
      if nargin==2
        shots = shots(1:nshots);
      end
    end
    
    function shots = get_silver_shots(table_filename,nshots)
      % get vector of shot numbers where silver shots are stored
      T = readtable(table_filename);
      shots = T.silver_shot;
      shots = shots(shots~=0);
      if nargin==2
        shots = shots(1:nshots);
      end
    end
  end
  
  methods(Static,Access=private)
    function success = write(source_shot,dest_shot,fort_equiv)
      fprintf('source shot: %d ',source_shot);
      fprintf('fort_equiv: %d ', fort_equiv);
      
      % Get static tree version
      mdsopen('static',-1);
      static = mdsvalue('\version[$1]',int32(source_shot));
      
      try
        % load shot parameters
        MG = mgp(source_shot);
        MG.itamax(:) = 100;
        MG.debug_p1 = sprintf('''static'',%d',static); % Must be a string
        if ~MG.check, MG.fix; end
        assert(MG.check,'mgp checks did not pass');
        % save to user shot number
        fprintf('destination shot: %d ',dest_shot);
        MG.save(dest_shot)
        
        % run fbt/mga ...
        if fort_equiv
          shotdesign_dir = fileparts(fileparts(mfilename('fullpath'))); % Path to the current shotdesign version
          % ... FORTRAN versions
          fprintf('run fbt.f .. ')
          fbtcmd = sprintf('%s/fbt %d > fbt.log 2>&1',shotdesign_dir,dest_shot);
          cleanupobj = onCleanup(@() delete('fbt.log'));
          assert(system(fbtcmd) == 0,'Error when running fbt');
          assert(system('grep "no convergence" fbt.log -q') == 1,'FBT did not converge for at least one run, cannot continue');
          assert(logical(exist('fbtmeas.dat','file')),'fbtmeas.dat missing, cannot continue');
          fprintf('run mga.f .. ')
          mgacmd = sprintf('%s/mga %d > mga.log 2>&1',shotdesign_dir,dest_shot);
          cleanupobj2 = onCleanup(@() delete('mga.log'));
          assert(system(mgacmd) == 0,'Error when running mga');
        else
          % ... MATLAB versions
          fprintf('run fbt.m .. ')
          [~,~,LY] = fbte(dest_shot);
          assert(all(LY.isconverged),'some FBT time slices did not converge')
          fprintf('run mga.m .. ');
          mga(dest_shot,false); % no output arguments to write to node, second argument selects FORTRAN compatibility mode
        end
        
        fprintf('done\n');
        success=true;
      catch ME
        fprintf('\n');
        warning('failed storing shot: error message was: \n %s',getReport(ME));
        success=false;
      end
    end
  end
end
