#!/usr/bin/env python

# --------------------------------------------------------- #
#  Simple python script to calculate prime numbers!
#   
#  Author: Troels C. Petersen (Niels Bohr Institute)
#  Date:   18-08-2014
#
#  Run by: ./CalcPrimeNumbers.py
# --------------------------------------------------------- #

from ROOT import *        # Import ROOT libraries (to use ROOT)
import math               # Import MATH libraries (for extended MATH functions)

Nmax = 10000              # Maximum prime number
N    = 2                  # Starting value

SavePlots = False         # Determining if plots are saved or not

gStyle.SetOptStat("emr")
gStyle.SetOptFit(1111)

hist_prime = TH1D("hist_prime", "Prime number distribution", 100, 0.0, Nmax)


# --------------------------------------------------------- #
# Calculate prime numbers:
# --------------------------------------------------------- #

while (N < Nmax):
    hit = 0                    # Define a variable "hit", with starting value 0.
    for i in range (2,int(sqrt(N))+1):
        if (N%i == 0):         # If the remainder of the integer division is 0, hit = 1
            hit = 1
    if (hit == 0):             # If no number gave a perfect integer division (i.e. with remainder 0)...
        hist_prime.Fill(N)     # the it is a prime number. Fill it into the histogram.
    N += 1                     # Increase N by one.

# --------------------------------------------------------- #
# Draw output:
# --------------------------------------------------------- #

canvas = TCanvas( "canvas", "canvas", 50, 50, 1200, 600 )

hist_prime.GetXaxis().SetRangeUser(0, Nmax)         # From the histogram, we get the x-axis, and set the range.
hist_prime.SetMinimum(0)                            # We also set the minimum.
hist_prime.GetXaxis().SetTitle("Prime number")      # We give the x-axis a label.
hist_prime.GetYaxis().SetTitle("Frequency")         # We give the y-axis a label.

hist_prime.SetMarkerColor(kBlue)                    # Set the histogram marker color to blue.
hist_prime.SetMarkerStyle(20)                       # Set the histogram marker style to filled circles.

fit_prime = TF1("fit_prime", "[0] + [1]/log(x)", 0, N)
fit_prime.SetParameters(0.0, 10.0)         # Set the starting values of [0] and [1] to 0 and 10.
fit_prime.SetLineColor(kRed)               # Set the line color to red.
fit_prime.SetLineWidth(4)                  # Set tne line width to 4.

hist_prime.Fit("fit_prime", "L")
hist_prime.Draw("e")

leg = TLegend( 0.25, 0.73, 0.60, 0.88 )
leg.SetFillColor(kWhite);                  # Let the fill color be white
leg.SetLineColor(kWhite);                  # Let the edge be white (i.e. no edge)
leg.AddEntry(hist_prime, " Prime number distribution", "PL")       # Histogram represented by a point "P"
leg.AddEntry(fit_prime,  " Prime number fit result", "L")         # Fit function represented by a line "L"
leg.Draw()                                 # Draw the legend

text = TLatex()                            # Useful for drawing text on plots
text.SetNDC()                              # Use "absolute" coordinates, i.e. [0,1]x[0,1]
text.SetTextFont(42)                       # One of the "standard" fonts
text.SetTextColor(1)
text.SetTextSize(0.05)
text.DrawLatex(0.15, 0.17, "Fit function: f(x) = C_{0} + #frac{C_{1}}{log(x)}")

if (SavePlots):
    canvas.SaveAs("PrimeNumberDistribution.pdf")              # Save plot (format follow extension name)



raw_input('Press Enter to exit')