%STOCHASTICFORCING class definition
%
%   Usage:
%      stochasticforcing=stochasticforcing();

classdef stochasticforcing
	properties (SetAccess=public)
		isstochasticforcing	= 0;
		fields					= NaN;
		defaultdimension		= 0;
		default_id				= NaN;
		covariance				= NaN;
		stochastictimestep   = 0;
		randomflag				= 1;
	end
	methods
		function self = stochasticforcing(varargin) % {{{
			switch nargin
				case 0
					self=setdefaultparameters(self);
				otherwise
					error('constructor not supported');
			end
		end % }}}
		function self = extrude(self,md) % {{{
			%Nothing for now
		end % }}}
		function self = setdefaultparameters(self) % {{{
			self.isstochasticforcing	= 0; %stochasticforcing is turned off by default
			self.randomflag				= 1; %true randomness is implemented by default
		end % }}}
		function md = checkconsistency(self,md,solution,analyses) % {{{

			%Early return
			if ~self.isstochasticforcing, return; end

			num_fields = numel(self.fields);
			if(self.stochastictimestep==0)
				md.stochasticforcing.stochastictimestep = md.timestepping.time_step; %by default: stochastictimestep set to ISSM time step
			end

			%Check that covariance matrix is positive definite
			try
				chol(self.covariance);
			catch
				error('md.stochasticforcing.covariance is not positive definite');
			end

			%Check that all fields agree with the corresponding md class and if any field needs the default params
			checkdefaults	= false; %need to check defaults only if one of the fields does not have its own dimensionality
			structstoch		= structstochforcing();
			for field=self.fields
				%Checking agreement of classes
				if(contains(field,'SMBautoregression'))
					mdname = structstoch.mdnames(find(strcmp(structstoch.fields,char(field))));
					if~(isequal(class(md.smb),char(mdname)))
						error('md.smb does not agree with stochasticforcing field %s', char(field));
					end
				end
				if(contains(field,'SMBforcing'))
					mdname = structstoch.mdnames(find(strcmp(structstoch.fields,char(field))));
					if~(isequal(class(md.smb),char(mdname)))
						error('md.smb does not agree with stochasticforcing field %s', char(field));
					end
				end
				if(contains(field,'FrontalForcings'))
					mdname = structstoch.mdnames(find(strcmp(structstoch.fields,char(field))));
					if~(isequal(class(md.frontalforcings),char(mdname)))
						error('md.frontalforcings does not agree with stochasticforcing field %s', char(field));
					end
				end
				if(contains(field,'Calving'))
					mdname = structstoch.mdnames(find(strcmp(structstoch.fields,char(field))));
					if~(isequal(class(md.calving),char(mdname)))
						error('md.calving does not agree with stochasticforcing field %s', char(field));
					end
				end
				if(contains(field,'BasalforcingsFloatingice'))
					mdname = structstoch.mdnames(find(strcmp(structstoch.fields,char(field))));
					if~(isequal(class(md.basalforcings),char(mdname)))
						error('md.basalforcings does not agree with stochasticforcing field %s', char(field));
					end
				end
				if(contains(field,'BasalforcingsSpatialDeepwaterMeltingRate'))
					mdname = structstoch.mdnames(find(strcmp(structstoch.fields,char(field))));
					if~(isequal(class(md.basalforcings),char(mdname)))
						error('md.basalforcings does not agree with stochasticforcing field %s', char(field));
					end
				end
				if(contains(field,'BasalforcingsDeepwaterMeltingRateAutoregression'))
					mdname = structstoch.mdnames(find(strcmp(structstoch.fields,char(field))));
					if~(isequal(class(md.basalforcings),char(mdname)))
						error('md.basalforcings does not agree with stochasticforcing field %s', char(field));
					end
				end
				if(contains(field,'WaterPressure'))
               mdname = structstoch.mdnames(find(strcmp(structstoch.fields,char(field))));
               if~(isequal(class(md.friction),char(mdname)))
                  error('stochasticforcing field %s is only implemented for default friction', char(field));
               end
               if(md.friction.coupling~=0 && md.friction.coupling~=1 && md.friction.coupling~=2)
                  error('stochasticforcing field %s is only implemented for cases md.friction.coupling 0 or 1 or 2', char(field));
               end
               if(any(md.friction.q==0))
                  error('stochasticforcing field %s requires non-zero q exponent',char(field));
               end
            end
				%Checking for specific dimensions
				if ~(strcmp(field,'SMBautoregression') || strcmp(field,'FrontalForcingsRignotAutoregression') || strcmp(field,'BasalforcingsDeepwaterMeltingRateAutoregression'))
					checkdefaults = true; %field with non-specific dimensionality
				end
			end

			%Retrieve all the field dimensionalities
			dimensions = self.defaultdimension*ones(1,num_fields);
			indSMBar   = -1; %about to check for index of SMBautoregression
			indTFar	  = -1; %about to check for index of FrontalForcingsRignotAutoregression
			indBDWar	  = -1; %about to check for index of BasalforcingsDeepwaterMeltingRateAutoregression
			if any(contains(self.fields,'SMBautoregression'))
				indSMBar = find(contains(self.fields,'SMBautoregression')); %index of SMBar, now check for consistency with other ar timesteps 
				dimensions(indSMBar) = md.smb.num_basins;
				if(md.smb.ar_timestep<self.stochastictimestep)
					error('SMBautoregression cannot have a timestep shorter than stochastictimestep');
				end
			end
			if any(contains(self.fields,'FrontalForcingsRignotAutoregression'))
				indTFar	= find(contains(self.fields,'FrontalForcingsRignotAutoregression')); %index of TFar, now check for consistency with other ar timesteps 
				dimensions(indTFar) = md.frontalforcings.num_basins;
				if(md.frontalforcings.ar_timestep<self.stochastictimestep)
					error('FrontalForcingsRignotAutoregression cannot have a timestep shorter than stochastictimestep');
				end
			end
			if any(contains(self.fields,'BasalforcingsDeepwaterMeltingRateAutoregression'))
				indBDWar	= find(contains(self.fields,'BasalforcingsDeepwaterMeltingRateAutoregression')); %index of BDWar, now check for consistency with other ar timesteps 
				dimensions(indBDWar) = md.basalforcings.num_basins;
				if(md.basalforcings.ar_timestep<self.stochastictimestep)
					error('BasalforcingsDeepwaterMeltingRateAutoregression cannot have a timestep shorter than stochastictimestep');
				end
			end
			size_tot = sum(dimensions);

			if(indSMBar~=-1 && indTFar~=-1) %both autoregressive models are used: check autoregressive time step consistency
				if(md.smb.ar_timestep~=md.frontalforcings.ar_timestep)
					crossentries = reshape(self.covariance(1+sum(dimensions(1:indSMBar-1)):sum(dimensions(1:indSMBar)),1+sum(dimensions(1:indTFar-1)):sum(dimensions(1:indTFar))),1,[]);
					if any(crossentries~=0)
						error('SMBautoregression and FrontalForcingsRignotAutoregression have different ar_timestep and non-zero covariance');
					end
				end
			end
			if(indSMBar~=-1 && indBDWar~=-1) %both autoregressive models are used: check autoregressive time step consistency
				if(md.smb.ar_timestep~=md.basalforcings.ar_timestep)
					crossentries = reshape(self.covariance(1+sum(dimensions(1:indSMBar-1)):sum(dimensions(1:indSMBar)),1+sum(dimensions(1:indBDWar-1)):sum(dimensions(1:indBDWar))),1,[]);
					if any(crossentries~=0)
						error('SMBautoregression and BasalforcingsDeepwaterMeltingRateAutoregression have different ar_timestep and non-zero covariance');
					end
				end
			end
			if(indTFar~=-1 && indBDWar~=-1) %both autoregressive models are used: check autoregressive time step consistency
				if(md.frontalforcings.ar_timestep~=md.basalforcings.ar_timestep)
					crossentries = reshape(self.covariance(1+sum(dimensions(1:indTFar-1)):sum(dimensions(1:indTFar)),1+sum(dimensions(1:indBDWar-1)):sum(dimensions(1:indBDWar))),1,[]);
					if any(crossentries~=0)
						error('FrontalForcingsRignotAutoregression and BasalforcingsDeepwaterMeltingRateAutoregression have different ar_timestep and non-zero covariance');
					end
				end
			end
			md = checkfield(md,'fieldname','stochasticforcing.isstochasticforcing','values',[0 1]);
			md = checkfield(md,'fieldname','stochasticforcing.fields','numel',num_fields,'cell',1,'values',supportedstochforcings());
			md = checkfield(md,'fieldname','stochasticforcing.covariance','NaN',1,'Inf',1,'size',[size_tot,size_tot]); %global covariance matrix
			md = checkfield(md,'fieldname','stochasticforcing.stochastictimestep','NaN',1,'Inf',1,'>=',md.timestepping.time_step);
			md = checkfield(md,'fieldname','stochasticforcing.randomflag','numel',[1],'values',[0 1]);
			if(checkdefaults) %need to check the defaults
				md = checkfield(md,'fieldname','stochasticforcing.defaultdimension','numel',1,'NaN',1,'Inf',1,'>',0);
				md = checkfield(md,'fieldname','stochasticforcing.default_id','Inf',1,'>=',0,'<=',self.defaultdimension,'size',[md.mesh.numberofelements,1]);
			end
		end % }}}
		function disp(self) % {{{
			disp(sprintf('   stochasticforcing parameters:'));
			fielddisplay(self,'isstochasticforcing','is stochasticity activated?');
			fielddisplay(self,'fields','fields with stochasticity applied, ex: [{''SMBautoregression''}], or [{''SMBforcing''},{''DefaultCalving''}]');
			fielddisplay(self,'defaultdimension','dimensionality of the noise terms (does not apply to fields with their specific dimension)');
			fielddisplay(self,'default_id','id of each element for partitioning of the noise terms (does not apply to fields with their specific partition)');
			fielddisplay(self,'covariance','covariance matrix for within- and between-fields covariance (units must be squared field units)');
			fielddisplay(self,'stochastictimestep','timestep at which new stochastic noise terms are generated (default: md.timestepping.time_step)');
			fielddisplay(self,'randomflag','whether to apply real randomness (true) or pseudo-randomness with fixed seed (false)');
			disp('Available fields:');
			for field=supportedstochforcings()
				fprintf('   %s\n',string(field));
			end
		end % }}}
		function marshall(self,prefix,md,fid) % {{{

			yts=md.constants.yts;
			num_fields = numel(self.fields);

			WriteData(fid,prefix,'object',self,'fieldname','isstochasticforcing','format','Boolean');
			if ~self.isstochasticforcing 
				return
			else

				if(self.stochastictimestep==0)
					disp('      stochasticforcing.stocahstictimestep not specified: set to md.timestepping.time_step');
					self.stochastictimestep = md.timestepping.time_step; %by default: stochastictimestep set to ISSM time step
				end

				%Retrieve dimensionality of each field
				dimensions = self.defaultdimension*ones(1,num_fields);
				ind = 1;
				for field=self.fields
					%Checking for specific dimensions
					if(strcmp(field,'SMBautoregression'))
						dimensions(ind) = md.smb.num_basins;
					end
					if(strcmp(field,'FrontalForcingsRignotAutoregression'))
						dimensions(ind) = md.frontalforcings.num_basins;
					end
					if(strcmp(field,'BasalforcingsDeepwaterMeltingRateAutoregression'))
						dimensions(ind) = md.basalforcings.num_basins;
					end
					ind = ind+1;
				end

				%Scaling covariance matrix (scale column-by-column and row-by-row)
				scaledfields = {'BasalforcingsDeepwaterMeltingRateAutoregression','BasalforcingsSpatialDeepwaterMeltingRate','DefaultCalving','FloatingMeltRate','SMBautoregression','SMBforcing'}; %list of fields that need scaling *1/yts
				tempcovariance = self.covariance; %copy of covariance to avoid writing back in member variable
				for i=1:num_fields
					if any(strcmp(scaledfields,self.fields(i)))
						inds = [1+sum(dimensions(1:i-1)):1:sum(dimensions(1:i))];
						for row=inds %scale rows corresponding to scaled field
							tempcovariance(row,:) = 1./yts*tempcovariance(row,:);
						end
						for col=inds %scale columns corresponding to scaled field
							tempcovariance(:,col) = 1./yts*tempcovariance(:,col);
						end
					end
				end
				%Set dummy default_id vector if defaults not used
				if isnan(self.default_id)
					self.default_id = zeros(md.mesh.numberofelements,1);
				end

				WriteData(fid,prefix,'data',num_fields,'name','md.stochasticforcing.num_fields','format','Integer');
				WriteData(fid,prefix,'object',self,'fieldname','fields','format','StringArray');
				WriteData(fid,prefix,'data',dimensions,'name','md.stochasticforcing.dimensions','format','IntMat');
				WriteData(fid,prefix,'object',self,'fieldname','default_id','data',self.default_id-1,'format','IntMat','mattype',2); %0-indexed
				WriteData(fid,prefix,'object',self,'fieldname','defaultdimension','format','Integer');
				WriteData(fid,prefix,'data',tempcovariance,'name','md.stochasticforcing.covariance','format','DoubleMat');
				WriteData(fid,prefix,'object',self,'fieldname','stochastictimestep','format','Double','scale',yts);
				WriteData(fid,prefix,'object',self,'fieldname','randomflag','format','Boolean');
			end
		end % }}}
	end
end
function list = supportedstochforcings() % {{{
	% Defines list of fields supported
	% by the class md.stochasticforcing

	list = structstochforcing();
	list = list.fields;
end % }}}
function structure = structstochforcing() % {{{
	% Defines structure with list of fields
	% supported and corresponding md names
	structure.fields = {...
		'BasalforcingsDeepwaterMeltingRateAutoregression',...
		'BasalforcingsSpatialDeepwaterMeltingRate',...
		'DefaultCalving',...
		'FloatingMeltRate',...
		'FrictionWaterPressure',...
		'FrontalForcingsRignotAutoregression',...
		'SMBautoregression',...
		'SMBforcing'
		};
	structure.mdnames = {...
		'autoregressionlinearbasalforcings',...
		'spatiallinearbasalforcings',...
		'calving',...
		'basalforcings',...
		'friction',...
		'frontalforcingsrignotautoregression',...
		'SMBautoregression',...
		'SMBforcing'
	};
end % }}}
