time_evolved_hist

from math import *
import numpy as np
import random
import matplotlib.pyplot as plt
import imageio.v3
import os

######### Here are parameters the user can tweak
dt = 1000   # Units: ??
de = 0.5    # Units: mEv, presumably
largest_time = 10000   # Manually set this
        # because most of what happens later on
        # is not very interesting

######## Here are constants the user shouldn't tweak
gap_energy = 0.1805
planck = 4.135667696 * 10**-15

######## Here are variables we fill later on
numsubgap = 0
tot_subgap_energy = 0.0
total_energy = 0.0
phononids = []
phononEnergies = []
phononTimes = []
phononEnergy = 0.0
phononTime = 0.0
largest_energy = 0.0
phononData = {}
energyDist = []
xlabel = []
ylabel = []

print("Finding data . . . ")
directory = "Data/"
path = directory + "StepData.txt"
f = open(path, "r")
lines = f.readlines()
f.close()

print("Parsing data. . . ")
numLines = 0
#Add energies for each event
for i in range(0, len(lines)):
    line = lines[i]
    if(line == "New Event\n"):
        break
    
    numLines += 1
    line_data = line[:-1].split(",")
    phononid = int(line_data[0])
    phononEnergy = float(line_data[1])
    phononTime = float(line_data[2])
    if(phononTime > largest_time):
        continue          # Only look at a certain time range

    if(line_data[3][0] != "p"):  # Only look at phonons!
        continue
    if(line_data[5] != "TargetL"):  # Only look at phonons
        continue                    # in silicon
    phononids.append(phononid)
    phononEnergies.append(phononEnergy)
    phononTimes.append(phononTime)

    # Find largest energy and time
    if(phononEnergy > largest_energy):
        largest_energy = phononEnergy
    total_energy += phononEnergy
    if(phononEnergy < 2 * gap_energy):
        numsubgap += 1
        tot_subgap_energy += phononEnergy
  
num_time_bins = int(largest_time / dt)
# num_time_bins = 1000
num_energy_bins = ceil(largest_energy / de)
setids = list(set(phononids))
setids.sort()
for i in range(len(setids)):
    phononData[setids[i]] = [[],[]]

for i in range(len(phononids)):
    phononData[phononids[i]][0].append(phononEnergies[i])
    phononData[phononids[i]][1].append(phononTimes[i])

for i in range(num_time_bins + 1):
    energyDist.append([])
    for j in range(num_energy_bins + 1):
        energyDist[i].append(0)

#iterate through all phonons
timea = 0.0
timeb = 0.0
timeaint = 0
timebint = 0
cur_energy = 0.0
energyind = 0
timeind = 0
pdata = []
lenp = 0

# iterate through IDs to see what time interval and energy interval the step belongs in
print("Computing dist of phonons per time...")
counter = 0  #To make sure we're not getting stuck
maxPhononNumInBin = 0    # Use this to restrict y axis
                          # of gif later on
for ID in phononData:
    pdata = phononData[ID]
    lenp = len(pdata[1])
    counter += 1
    if (counter % 1000 == 0): print(counter)
    # iterates through all steps for each phonon
    # finds the slices of time and the energy that the phonon had at those times and
    #   increases the counter for that grid location
    for i in range(lenp - 1):
        timea = pdata[1][i]
        timeb = pdata[1][i + 1]
        timeaint = floor(timea / dt) * dt
        timebint = ceil(timeb / dt) * dt
        cur_energy = pdata[0][i]
        energyind = num_energy_bins - int(cur_energy / de)
      
        # if(timeaint == 0):
        #     energyDist[0][energyind] += 1
        #     # energyDist[0][energyind] += cur_energy
        for j in range(timeaint, timebint, dt):
            timeind = int(float(j) / dt)
            energyDist[timeind][energyind] += 1

            # recalc max number of phonons in any one bin
            if energyDist[timeind][energyind]>maxPhononNumInBin:
              maxPhononNumInBin = energyDist[timeind][energyind]
            # if(timeind < num_time_bins):
            #     energyDist[timeind][energyind] += 1
                # energyDist[timeind][energyind] += cur_energy

# tot_init_energy = 0.0
# for i in range(len(energyDist[0])):
#     tot_init_energy += energyDist[400][i]


# Find the average energy at each time. 
# is the curve really moving to the right?
'''
plt.scatter(phononTimes, phononEnergies)
plt.show()

x=[]  # Time
y=[]  # Avg phonon energy
for i in range(1,num_time_bins):
  tot_energy = 0
  num_phonons_this_time = 0
  for j in range(num_energy_bins):
    num_phonons_this_energy = energyDist[i][j]
    num_phonons_this_time += num_phonons_this_energy
    tot_energy += num_phonons_this_energy*(num_energy_bins-j)*de
    
  av_energy = tot_energy/num_phonons_this_time if num_phonons_this_time>0 else 0
  x.append(i)
  y.append(av_energy)
plt.scatter(x, y)
plt.show()
'''

tot_init_energy = energyDist[1][0]
# print("Total Initial Energy: " + str(0.001 * tot_init_energy))


print("Plotting . . . ")
print("Proportion of subgap phonons: " + str(float(numsubgap) / numLines))
print("Total Subgap Energy: " + str(tot_subgap_energy) + ", Total Energy: " + str(total_energy))
print("Subgap Energy Fraction: " + str(float(tot_subgap_energy) / total_energy))

####################################################################
# Plot phonon frequency histograms per energy at each time point.
# Using Mira's pulse_data_analysis as template, turn this into a gif.

start_time_interval=0
frames=num_time_bins
framename=[]
for index in range(start_time_interval,start_time_interval+frames):

    ########################### First, create energy hist at start_time_interval.
    fig1 = plt.figure()                 # Figure creation
    ax1 = fig1.add_subplot(1, 1, 1)
    ax1.set_xlim(0,(largest_energy+1)/ (10**3) / planck / (10**12))
    ax1.set_ylim(0,maxPhononNumInBin+1)
    ax1.set_title("Time Evolution of Phonons in TargetL")
    energy_bin_x_data = np.arange(largest_energy+de, 0, -de) #Get energy bins
    energy_bin_x_data = energy_bin_x_data / (10**3) / planck / (10**12)
    frequency_phonons_y_data = energyDist[index]    
                                        # and freq for each bin at desired time

    ax1.plot(energy_bin_x_data, frequency_phonons_y_data)
    ax1.set_xlabel('Frequency (THz)')
    ax1.set_ylabel('Number of Phonons')

    ########################### And save it to a "temp" folder.
    if (not os.path.isdir("temp")): # Make temp folder if it doesn't 
        os.mkdir("temp")            # already exist
                                                    
    name="temp/pic"+str(index)+".png"
    fig1.savefig(name,dpi=150,facecolor='white',transparent=False,bbox_inches='tight')
    framename.append(name)
    plt.close()

################################ Second, make a gif of the images in the temp folder    
images = []
for filename in framename:
    images.append(imageio.imread(filename))
duration = 0.5
imageio.mimwrite("phonon_energy_hist"+".gif", images, duration=duration,
                 loop=0)