clear
set(0,'DefaultAxesFontSize',25)
set(0,'DefaultLinelinewidth',1)

% DESCRIPTION ::::::::::::::::::::
% A matlab routine to fit bolometric SN
% light curves with magnetar central engine model (with properties
% as in Kasen and Bildsten 2010) following the Arnett 1982
% formalism, with modification of kinetic energy treatment as described in
% Inserra, Smartt, Jerkstrand et al 2013.
% Author : A. Jerkstrand.
% Article reference : Inserra, Smartt, Jerkstrand et al ApJ 2013
% (denoted I13 below) 
% :::::::::::::::::::::::::::::::::

% CONTROLS ::::::::::::::::::::::::
obj = 'test'   % Select object to fit (4 char string).
kappa = 0.1;   % Opacity for UVOIR radiation (cm^2/g) 
alpha = 10;    % Power law exponent for envelope density profile
uselims = 0;   % 0: Ignore upper limits in fits. 1: Include upper limits in fits 
dotrap = 1;    % 0: Assume full X-ray/gamma trapping 1: Compute trapping.
kappa_pp = 0.01; % Opacity to magnetar radiation (cm^2/g), only used if dotrap = 1
flow = 0.1;    % Minimum multiplication factor from starting guess to probe
fhigh = 10;    % Maximum multiplication factor from starting guess to probe
N_Mej = 5;     % Number of bins in Mej for coarse scan
N_B14 = 5;     % Number of bins in B14 for coarse scan
N_Pms = 5;     % Number of bins in Pms for coarse scan
Nshift = 5;    % Number of bins in time shift for coarse scan 
N_nimass = 5;

Nevals = 2000  % Parameter for fminsearch
tol = 1E-10    % Parameter for fminsearch

t_end = 300;   % Ending time of plotted light curve in days
save = 0;      % Save figures?

% -----------------------
% Usage : 
% Step 1 : Create a data input file along the following format:
%   First 2 lines : comments. Then
%   Column 1: Observer-frame epoch relative to peak (days)
%   Column 2: Logarithm of SN Rest-frame Bolometric luminosity (log10 L_bol (erg/s)) 
%   Column 3: Error on log10 L_bol
%   Any upper limits should be listed at the end, and controlled by
%   parameter "skip" below (=number of upper limits)
%
% Step 2: Copy and paste the section ("if (obj == 'test')") below
%   , changing the various variables to the new data set
%
% Step 3 : Set the fit variables above
%
% Step 4 : Run the code. Observed velocities and temperatures 
%   need to be added into Figs 2 and 3.
%
% Note that in the Arnett 1982 formalism, the light curve depends
% only on the power source and on the effective diffusion time
% scale tau_m = 1.05/sqrt(13.7*c)*kappa^(1/2)*Mej^(3/4)*Ek^(-1/4) s
% for constant density. The difference between the 56Ni/56Co case
% (as described in Arnett 1982) and the magnetar case is
% that energy input by the magnetar can be so large that the
% acceleration of the ejecta must be taken into account. We found
% by calibrating the solutions to the radiation hydrodynamic models
% in Kasen and Bildsten 2010 that using Ek = E0 + 1/2(Emag - Erad)
% gives excellent agreement : here E0 is the explosion energy, Emag
% is the total energy input by the magnetar, and Erad is the total
% radiated SN luminosity. In the fitting process here two steps are
% needed to get the final M_ej and velocities/temperatures correct:
% make first an initial fit with any value for E_51 (say 1) :
% obtain Emag and Erad from this run and plug into the expression
% for E_51 - rerun and the final M_ej value and
% temperature/velocity plots are obtained.

% Data and starting guesses for each SN:

if (obj == 'test')
  fid = fopen('test.txt'); % Input file
  shiftmin = 0;  % minimum time shift allowed from first epoch
  shiftmax = 50; % maximum time shift allowed
  skip = 0;   % Number of upper limits (in end of input file) 
  redshift = 0.100; % redshift
  if (dotrap == 0) % If assuming full trapping
      E_51 = 1 + 0.5*(0.31 - 0.17);  % See description above : Take numbers from Emag and
                                     % Erad after first run (they
                                     % do not depend on E_51 in the
                                     % fully trapped case)
      Mej_mag = 3;
      B14_mag = 3.6; 
      Pms_mag = 7.2; 
      Mni_mag = 0.1;
      savename = 'magfit_test';
     
  else         
      E_51 = 1 + 0.5*(0.22 - 0.17);  % Replace by 1 + 0.5*(Emag - Etot), with Emag and Etot taken from iterative simulations. In the non-fully trapped case, degeneracy between M^(3/4)E^(-1/4) is broken, so iteration is needed.
      Mej_mag = 1;
      B14_mag = 2.1; 
      Pms_mag = 6.4;
      Mni_mag = 0.1;
      savename = 'magfit_test_trap';
      
  end 
  name = 'Test'
  save = 0;
  % For plotting:
  y0 = 40; % yrange-min
  y1 = 45; % yrange-max
  t1 = 43.8; t2= 43.5; t3=43.2; t4=42.9; t5=42.6 % text positions
  xlv = 50;

end

% ADD MORE DATA HERE AS FOR EXAMPLE ABOVE...

% if (obj = 'xxx')
%    ...

% END OF USER INPUT ::::::::::::::::::::::::::::::::::::::::::::

forestgreen = [34 139 34]/255.;
gold = [0.75 0.75 0];

% READ IN DATA ------------------ :

fgetl(fid); % Comments line 1

fgetl(fid); % Comments line 2 

Z = fscanf(fid,'%f',[3 Inf]); % Data, 3 columns (see above)

fclose(fid);

tvec_days = Z(1,1:length(Z)-skip);  % Col 1 : Time in days observer frame
                                    % (arbitrary zero time),
                                    % Skip upper limits

ydata_log = Z(2,1:length(Z)-skip);  % Col 2 : Log10 L (erg/s)
                                    % bolometric
                                    % SN-frame lum. Skip
                                    % upper limits

sigma_log = Z(3,1:length(Z)-skip);  % Errors in log10 L_bol

N_data = length(tvec_days);  % Number of epochs

if (skip > 0)  % If there are any upper limits in the data (they
               % are always at the end)...

  tvec_skip = Z(1,length(Z)-skip+1:length(Z)-skip+1+(skip-1));   % Epochs for the upper limits...

  ydata_skip = Z(2,length(Z)-skip+1:length(Z)-skip+1+(skip-1));  % log10 L for the upper limits...

  sigma_skip = Z(3,length(Z)-skip+1:length(Z)-skip+1+(skip-1));   % Errors in log L (dont include BC)

else % There are no upper limits

   tvec_skip = [];

   ydata_skip = [];

end

% ROUGH SPACE SCAN (exact best fit found by later optimization from this starting point) ------------------------

Mej_min = Mej_mag*flow;  % Minimum ejecta mass in course search
Mej_max = Mej_mag*fhigh;  % Maximum ejecta mass in course search

B14_min = B14_mag*flow;  % Minimum B field in 10^14 G in course search
B14_max = B14_mag*fhigh; % Maximum B field in 10^14 G in course search

Pms_min = Pms_mag*flow;  % Minimum period in ms in course search
Pms_max = Pms_mag*fhigh; % Maximum period in ms in course search

Mni_min = Mni_mag*flow; % Minimum 56Ni mass in course search
Mni_max = Mni_mag*fhigh; % Maximum 56Ni mass in course search

% shiftmin and shiftmax are set above

Mej_vector = linspace(Mej_min, Mej_max, N_Mej);
B14_vector = linspace(B14_min, B14_max, N_B14);
Pms_vector = linspace(Pms_min, Pms_max, N_Pms);
shift_vector = linspace(shiftmin, shiftmax, Nshift);
Mni_vector = linspace(Mni_min, Mni_max, N_nimass);

bestchi2 = 1e50;  % Initialize best chi2 to a large value...

indb = find(tvec_days < 500); % TEMP

for i = 1:N_Mej  % For each ejecta mass..

    for j = 1:N_B14  % For each magnetic field value...

        for k = 1:N_Pms   % For each spin period..

            for l = 1:Nshift   % For each time shift...

		 for m = 1:N_nimass % For each 56Ni mass...
  
                   % Extract parameter values:
                   x(1) = Mej_vector(i);
                   x(2) = B14_vector(j);
                   x(3) = Pms_vector(k);
                   x(4) = shift_vector(l);
                   x(5) = Mni_vector(m);
                 
                   if ((x(3) > 1.0) && (x(5) < x(1))) % Only try Pms > 1 ms periods, and M(56Ni) < M_ej

			chi2(i,j,k,l,m) = chi2function(x, E_51, kappa, kappa_pp, length(indb), tvec_days(indb), ydata_log(indb), sigma_log(indb), uselims, skip, tvec_skip, ydata_skip, redshift, dotrap);  % Compute the reduced chi2 value of this model
                                    
			if (chi2(i,j,k,l,m) < bestchi2)  % If new
                                                   % solution is
                                                   % the best one
                                                   % so far, store it...

		        bestchi2 = chi2(i,j,k,l,m);
                        bestMej = Mej_vector(i);
                        bestB14 = B14_vector(j);
                        bestPms = Pms_vector(k);
                        bestshift = shift_vector(l);
                        bestnimass = Mni_vector(m);

                       end

                    end
                    
                end

            end

        end

    end

end

bestMej_coarse = bestMej;
bestB14_coarse = bestB14;
bestPms_coarse = bestPms;
bestshift_coarse = bestshift;
bestnimass_coarse = bestnimass;
bestchi2_coarse = bestchi2;

disp('Best solution from coarse search: M_ej, B14, Pms, shift, chi2(reduced):')
bestMej
bestB14
bestPms
bestshift
bestnimass
bestchi2

% Now perform an exact determination of the chi2 minimum: (cannot
% do globally from scratch, search often diverges);

x0 = [bestMej, bestB14, bestPms, bestshift, bestnimass];  % Starting point is
                                              % best solution from
                                              % coarse search..

options = optimset('MaxFunEvals',Nevals);   % Increase maximum number of function evaluations, needed for some searches.
options = optimset('TolFun',tol);  % If upper limits are used,
                                   % this has to be small (the
                                   % default is 1e-4 and
                                   % optimization may terminate
                                   % prematurely for that value)

fval = 1E30
% Do constrained minimization of chi2:
  [x, fval] = fminsearchbnd(@(x) chi2function(x, E_51, kappa, kappa_pp, length(indb), tvec_days(indb), ydata_log(indb), sigma_log(indb),uselims, skip, tvec_skip, ydata_skip, redshift, dotrap), x0, [0 0 1.0 -Inf 0], [100 100 Inf Inf Inf], options) % Dont allow Pms below 1 ms. THe two bracket [] are min and max values allowed for each parameter
			
if (fval < bestchi2) % Was a new best fit found?
 % Refined best fit values:
 bestMej = x(1);    
 bestB14 = x(2);
 bestPms = x(3);
 bestshift = x(4);
 bestnimass = x(5);
 bestchi2 = fval;
 disp('Best solution exact search: M_ej, B14, Pms, shift, chi2(reduced):')
 bestMej
 bestB14
 bestPms
 bestshift
 bestnimass
 bestchi2
 Bestexact = x;

else

 disp('Refined search failed to improve solution.')

end

%bestMej = 2.2; % TEMP
%bestB14 = 1; % TEMP
%bestPms = 1000.0 % TEMP
%E_51 = 1 + 0*0.5*(4.62 - 0.93) % TEMP
%bestnimass = 0.08;
%kappa = 0.1 % TEMP
%dotrap = 1;

% PLOTTING ----------

% Constants
c       = 2.998e10; % Speed of light (cm/s)   
e_ni0   = 3.90e10;  % Energy release in 56Ni decay (erg/s/g)  CHECKED. Note Arnett 1982 has too high value 4.8e10
e_co0   = 6.78e9;   % Energy release in 56Co decay (erg/s/g)  CHECKED 
m_sol   = 2e33;     % Solar mass (g) 
tau_ni  = 6.077*24*60*60/log(2); % Decay time (half life / ln2) for
                                 % Ni56 (s) 6.077 --> 8.8 d CHECKED
                                 % note ln in matlab is log
tau_co  = 77.27*24*60*60/log(2); % Decay time (half life / ln2) for Co56 (s) 77.27d --> 111.5 d  CHECKED
kappa_gamma = 0.03;  % Effective opacity to gamma rays from radioactivity CHECKED

tau_m = 1.05*((kappa/(13.7*c))^0.5) * (((((bestMej*m_sol)^3))/(E_51*1e51))^0.25);  % Diffusion time (Arnett 1982 Eq 18, 19, 22, 23)  CHECKED

y = tau_m/(2*tau_ni);   % Arnett 1982 Eq 33 CHECKED

yp = tau_m/(2*tau_co);    % Arnett 1982 Eq 33, modified to 56Co
                          % decay CHECKED
 
N = 10000;    % Number of time points to compute light curve at, and also numer of integration steps for each time. TEMP This should be >~ 10000 to be accurate at late times

if (skip == 0) 

  tend = max(tvec_days(length(tvec_days)))

else
     
  tend = max(tvec_days(length(tvec_days)),tvec_skip(length(tvec_skip)))
      
end

tend = tend/(1+redshift)+bestshift
     
tvec_model = linspace(0.001,tend,N);   % Create time vector

Erad = 0;   % Initialize total emitted energy by SN 

Emag = 0;   % Initialize total emitted energy by magnetar

V_ej = sqrt(10./3.*E_51*1e51/(bestMej*2e33));   % Maximum ejecta
                                                % velocity,
                                                % assuming constant
                                                % density  % Arnett
                                                % 1980 after Eq 32,
                                                % also I13
                                                % Eq D11 CHECKED

for i = 1:N     % Loop over time steps..for each time the solution is an integral over previous time

    t = tvec_model(i)*24*3600;   % Time in seconds  %  CHECKED

    if (i == 1)

        dt = (tvec_model(2)-tvec_model(1))*24*3600;  % Time step in
                                                 % seconds TEMP
                                                 % CHECKED
    else
        
        dt = (tvec_model(i)-tvec_model(i-1))*24*3600;  % Time step in
                                                 % seconds TEMP
                                                 % CHECKED
    end
    
    % Calculate parameters for photospheric phase

    x = t/tau_m;   % Arnett 1982 Eq 32  CHECKED

    z = linspace(0.0001,x,N);    % Define limits of integration for A(z)  CHECKED

    Ep = 1.97e52*bestPms^(-2);   % CHECKED

    tp = 0.013*bestB14^(-2)*bestPms^2*365*24*3600;  % CHECKED

    if (dotrap == 1) % Compute magnetar radiation trapping?

        % Compute gamma ray deposition

        R = V_ej*z*tau_m;     % Radius vector...z*tau_m = time;  CHECKED

        rho = bestMej*2e33./(4*pi/3*R.^3); % CHECKED

        tau_56co = kappa_gamma*rho.*R;

        tau_pp = kappa_pp*rho.*R;  % CHECKED

        G_56Co = tau_56co./(tau_56co + 1.6);

        G_pp = tau_pp./(tau_pp + 1.6);    % Arnett 1982 Eq 51 CHECKED

        D_gamma_56Co =  G_56Co.*(1 + 2*G_56Co.*(1-G_56Co).*(1-0.75*G_56Co));    % Arnett 1982

        D_gamma_pp = G_pp.*(1 + 2*G_pp.*(1-G_pp).*(1-0.75*G_pp));    % Arnett 1982
                                                      % eq 50
                                                      % CHECKED
    else % Assume full trapping

        D_gamma_56Co = 1.;

        D_gamma_pp = 1.;

    end

    integrand_magnetar = Ep/tp*(2-1)./(1. + z*tau_m/tp).^2.*D_gamma_pp*exp(-x^2).*exp(z.^2).*2.*z;  % CHECKED 
          
    % 56Ni part: -------------------------------
  
    integrand1= exp(-x^2)*2*z.*exp((-2*z*y)+(z.^2)).*D_gamma_56Co;    % Define function A(z) Arnett 1982 Eq 31  CHECKED

    Lambda1 = sum(integrand1.*(x/N));                 % Numerically integrate A(z)  Results checked against table 1 in Arnett 1982 CHECKED
  
    % 56Co part: --------------------------------
   
    integrand2 = exp(-x^2)*2*z.*(exp(-2*z*yp) - exp(-2*z*y))/(1 - tau_ni/tau_co).*exp(z.^2).*(0.966*D_gamma_56Co + 0.034);   % Define function B(z) (extension of Eq 31 in Arnett 1982)  CHECKED
 
    Lambda2 = sum(integrand2.*(x/N));                 % Numerically integrate B(z)

    L_ph_1(i) = sum(integrand_magnetar.*(x/N)) + bestnimass*m_sol*(e_ni0*Lambda1 + e_co0*Lambda2);  % Bolometric TEMP
                                                 % luminosity at
                                                 % this time CHECKED
   
    Emag = Emag +  Ep/tp*(2-1)./(1. + t/tp).^2*dt*D_gamma_pp(length(D_gamma_pp));  % Total magnetar energy output so far..  CHECKED
    
    Erad = Erad + L_ph_1(i)*dt;   % Total SN luminosity emitted so far...  CHECKED
    
end

% Plot lightcurve ------------------------------

figure(1), clf, hold on, box on

  errorbar(tvec_days/(1.+redshift)+bestshift, ydata_log, sigma_log,'xr') % Convert time from observer-frame time to SN rest frame time

  % Plot upper limits, if there are any:
  if (skip > 0)

      plot(tvec_skip/(1.+redshift)+bestshift, ydata_skip, 'rv','Markersize',10)

  end

  % Plot the best-fitting model:
  plot(tvec_model, log10(L_ph_1),'k')  

  xlabel('Days since explosion');      

  ylabel('Log L (erg s^{-1})');

  xmax = tend*1.1
  %xmax=100 % TEMP
  xlim([-50 xmax]) % TEMP

  ylim([y0 y1])

  set(gca,'Xminortick','on')

  set(gca,'Yminortick','on')

  fs = 18

  xl = 145

  text(xl,t1,name,'Fontsize',fs+5,'Color','r')
  text(xl,t2,['M_{ej} = ' num2str(bestMej,2) ' M_{sun}'],'Fontsize',fs)
  text(xl,t3,['B = ' num2str(bestB14,2) '*10^{14} G'],'Fontsize',fs)
  text(xl,t4,['P = ' num2str(bestPms,2) ' ms'],'Fontsize',fs)
  text(xl,t5,['M(56Ni) = ' num2str(bestnimass,2) ' Msun'],'Fontsize',fs)

  if (save == 1)

      hgexport(gcf,savename, hgexport('factorystyle'),'Format','eps')

  end

% end of figure 1

% Compute photospheric velocities and temperatures:

for i = 1:N % For each time...

    t = tvec_model(i)*24*3600;   % Time in seconds CHECKED

    Rcore = V_ej*t;    % Outer ejecta radius in cm CHECKED

    rhocore = bestMej*2e33/(4*pi/3.*Rcore^3);  % Core density in
                                               % g/cm^3 I13
                                               % Eq D12 CHECKED
 
    taucore = kappa*Rcore*rhocore;   % Core optical depth in
                                     % optical CHECKED

    tauatm = taucore/(alpha-1.);   % Atmosphere optical depth in
                                   % optical CHECKED

    if (tauatm > 1)  % If atmosphere is still optically thick...

      Rstar = Rcore*((alpha-1.)/taucore)^(1./(1.-alpha));  % Eq D19
                                                           % in I13
                                                           % CHECKED 

      Tstar(i) = (L_ph_1(i)/(4*pi*Rstar^2*5.67e-5))^0.25;  % Effective temperature CHECKED

      Vstar(i) = Rstar/t;   % Photospheric velocity
				
    elseif ((tauatm + taucore) > 2)   % Weird solutions when core is almost thin, unreliable solutions for such late phases

       taumissing = 1 - tauatm;  % The optical depth that the core (between R* and Rcore) must provide

       Rstar = Rcore - taumissing/(kappa*rhocore);  % I13 Eq D20 CHECKED

       if (Rstar > 0.5*Rcore)   % Compute only if more than half the core is covered ny photosphere

          Tstar(i) = (L_ph_1(i)/(4*pi*Rstar^2*5.67e-5))^0.25;  % Blackbody formula CHECKED

          Vstar(i) = Rstar/t; % CHECKED

       else

	  Tstar(i) = 0;  % Stop when more than half core has been exposed

          Vstar(i) = 0;

       end

    else  % Weird solutions when core is almost thin, unreliable solutions for such late phases

        Tstar(i) = 0;

        Vstar(i) = 0;

    end

end

% Write to file

fid = fopen('magnetarfit.txt','w')

 for i = 1:length(tvec_model)

    fprintf(fid,'%12.8e %12.8e %12.8e %12.8e\n', tvec_model(i), L_ph_1(i), Vstar(i), Tstar(i))

 end

fclose(fid)

% PLOT EFFECTIVE TEMPERATURES -------------------------

figure(2), clf

  ind = find(Tstar > 0);

  plot(tvec_model(ind), Tstar(ind),'k')

  set(gca,'Yminortick','on')

  title('T_{eff}')

  % Add plotting of observed temperatures here... 

  xlabel('Days since explosion')

  set(gca,'Xminortick','on')

  ylim([0 2e4])

  if (save == 1)

    hgexport(gcf,strcat(savename,'_T'), hgexport('factorystyle'),'Format','eps')

  end

% PLOT PHOTOSPHERIC VELOCITIES----------------------

figure(3), clf, hold on, box on

  ind = find(Vstar > 0);

  plot(tvec_model(ind), Vstar(ind)/1e5/1e4,'k')
  
  fs = 15
  
  %errorbar(timetemp, vfeii/1e4,errorfeii/1e4, 'ok','Markersize',fs,'Markeredgecolor','k','Markerfacecolor',forestgreen)
  %errorbar(timetemp, vmgii/1e4,errormgii/1e4, 'xk','Markersize',fs,'Markerfacecolor',forestgreen,'Markeredgecolor',forestgreen)
  %errorbar(timetemp, voi/1e4,erroroi/1e4, 'dk','Markersize',fs,'Markerfacecolor',forestgreen,'Markeredgecolor','k')
  %fs = 15;
  %text(xlv,1.7e4/1e4,name,'Fontsize',30,'Color',forestgreen)

  xlim([-10 150])


  xlabel('Rest frame days')

  ylabel('V ($10^4$ km s$^{-1})$','Interpreter','latex')

  set(gca,'Xminortick','on')

  set(gca,'Yminortick','on')

  ylim([0 2e4/1e4])

  if (save == 1)

       hgexport(gcf,strcat(savename,'_v'), hgexport('factorystyle'),'Format','eps')

  end

      
