from __future__ import division # Make 1/2 return 0.5

# Copyright 2005 Matthew Dixon Cowles <matt@mondoinfo.com>

# Distributable under the Gnu General Public License.

# There are any number of ways to make this faster at the
# expense of making it at least slightly more obscure. It
# runs quite fast enough for me.

import datetime
from elementtree import ElementTree
import itertools
import operator
import os
import random
import time

kAddAsNewDays=60
kAlbumStdDevsToGet=6
kEncodeChars=("&","<",">") # Can't just do kEncodeCharsEncoded.keys(), must have order to do & first
kEncodeCharsEncoded={"&": "&amp;", "<": "&lt;", ">": "&gt;"}
kEncodeValues=("Artist","Album","Name")
kGenresToSwitchAmong=("Classical",None) # None for all others
kOutputFileDir=os.path.expanduser("~/Desktop")
kOutputFilename="pl-auto-generated.xml"
kPLMaxBytes=4*1000000000 # 4 "GB"
kPLNameTemplate="Auto-generated %Y-%m-%d %H:%M"
kTrackStdDevsToGet=4


class utc(datetime.tzinfo):
  def utcoffset(self,dtIgnored):
    return datetime.timedelta(0) 
  def dst(self,dtIgnored):
    return datetime.timedelta(0)
  def tzname(self,dt):
    return "UTC"
    
class localTimezone(datetime.tzinfo):
  def __init__(self):
    self.stdOffset = datetime.timedelta(seconds = -time.timezone)
    # time.daylight indicates presence of DST, not that it's currently DST
    if time.daylight:
      self.dstOffset = datetime.timedelta(seconds = -time.altzone)
    else:
      self.dstOffset = self.stdOffset

    self.dstDiff = self.dstOffset - self.stdOffset
    return None
    
  def utcoffset(self, dt):
    if self.isdst(dt):
      return self.dstOffset
    else:
      return self.stdOffset

  def dst(self, dt):
    if self.isdst(dt):
      return self.dstDiff
    else:
      return datetime.timedelta(0)

  def tzname(self, dt):
    return time.tzname[self.isdst(dt)]

  def isdst(self, dt):
    tt = (dt.year, dt.month, dt.day,
        dt.hour, dt.minute, dt.second,
        dt.weekday(), 0, -1)
    stamp = time.mktime(tt)
    tt = time.localtime(stamp)
    return tt.tm_isdst > 0


# Some people may think that this is a cheesy
# way to implement control flow. I don't think
# I'm among them. After all, it's not that
# much different from StopIteration.
class plFull(Exception):
  pass
  
def tryAdd(playList,tracks,which):
  totTrackSize=0
  for track in tracks:
    totTrackSize+=track["Size"]
  if totTrackSize+playList["Size"]>kPLMaxBytes:
    print "Not added, would make playlist too big"
    raise plFull
  else:
    for track in tracks:
      playList["All Track IDs"].append(track["Track ID"])
      playList[which].append(track["Track ID"])
    playList["Size"]+=totTrackSize
    
  print "So far %i tracks, %sB" % (len(playList["All Track IDs"]),qtyToReadable(playList["Size"]))
  return None
  

def qtyToReadable(val):
  scaleSuffixes=("","K","M","G","T","P")
  scale=0
  while val>=1024 and scale<len(scaleSuffixes)-1:
    val=val/1024.0
    scale+=1
  if val==int(val):
    decimals=0
  elif 10*val==int(10*val):
    decimals=1
  else:
    decimals=2
  return "%.*f%s" % (decimals,val,scaleSuffixes[scale])


# Swiped from random.py from the standard library. Create
# a copy of the list that's kind of
# randomly ordered. Elements are very likely to be
# reordered, but early items tend to stay early. This is
# probably not the best way to do it, but I don't know a
# better one.

def biasedSample(population):
  n=len(population)
  numToReturn=n
  result = [None] * numToReturn
  pool=population[:]
  for i in xrange(numToReturn):
    j = int(random.betavariate(1,2) * (n-i))
    result[i] = pool[j]
    pool[j] = pool[n-i-1]
  return result


def someStats(numSeq):
  stats={}
  mySeq=numSeq[:]
  mySeq.sort()
  stats["min"]=mySeq[0]
  stats["max"]=mySeq[-1]
  
  n=len(numSeq)
  stats["n"]=n
  mean=sum(numSeq)/n
  stats["mean"]=mean
  
  accum=0
  for num in numSeq:
    accum+=(num-mean)**2
  accum/=n
  stats["stdDev"]=accum**0.5
  
  aboutHalfway=int(n/2)-1
  if n % 2 ==0:
    stats["median"]=(mySeq[aboutHalfway]+mySeq[aboutHalfway+1])/2
  else:
    stats["median"]=mySeq[aboutHalfway]
  
  return stats
  
  
def median(numSeq):
  mySeq=numSeq[:]
  mySeq.sort()
  n=len(mySeq)
  aboutHalfway=int(n/2)-1
  if n % 2 ==0:
    return int(round((mySeq[aboutHalfway]+mySeq[aboutHalfway+1])/2))
  else:
    return mySeq[aboutHalfway]
  
  
def makeDate(utcStr):
  assert utcStr[-1]=="Z"
  t=time.strptime(utcStr,"%Y-%m-%dT%H:%M:%SZ")
  return datetime.datetime(t[0],t[1],t[2],t[3],t[4],t[5],0,utc())
  
  
def dictify(el):
  assert el.tag=="dict"
  eList=el.getchildren()
  assert len(eList) % 2 == 0
  d={}
  for ind in range(0,len(eList),2):
    kEl=eList[ind]
    vEl=eList[ind+1]
    assert kEl.tag=="key"
    
    if vEl.tag=="integer":
      v=int(vEl.text)
    elif vEl.tag=="string":
      v=vEl.text
    elif vEl.tag=="date":
      v=makeDate(vEl.text)
    elif vEl.tag=="true":
      v=True
    elif vEl.tag=="dict": 
      v=vEl # We'll dictify it later
    elif vEl.tag=="array":
      v=vEl.getchildren()
    else:
      assert 1==0, "Can't happen, tag is %s" % vEl.tag
    d[kEl.text]=v
  return d
  
  
  
def makePLText(tracks,playList):
  # Could prolly do this with ElementTree pretty easily, but
  # a few templates is prolly even easier.
  
  # Hmmm, maybe that wasn't as trivial as I had thought.
  
  plOrder=["Recently Added","Never Played","Most Played"]
  for genre in kGenresToSwitchAmong:
    if genre==None:
      plOrder.append("Other Genres")
    else:
      plOrder.append(genre)
  
  
  filePrologue="""<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple Computer//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0">
<dict>
<key>Major Version</key><integer>1</integer>
<key>Minor Version</key><integer>1</integer>
<key>Application Version</key><string>5.0</string>
<key>Tracks</key>
<dict>
"""

  fileTrackTemplateIDName="""<key>%s</key>
<dict>
<key>Track ID</key><integer>%s</integer>
<key>Name</key><string>%s</string>
"""

  fileTrackTemplateArtist="""<key>Artist</key><string>%s</string>
"""

  fileTrackTemplateAlbum="""<key>Album</key><string>%s</string>
"""

  fileTrackEnd="""</dict>
"""
  
  fileMiddleTemplate="""</dict>
<key>Playlists</key>
<array>
"""

  filePLStartTemplate="""<dict>
<key>Name</key><string>%s</string>
<key>All Items</key><true/>
<key>Playlist Items</key>
<array>
"""

  filePLTrackTemplate="""<dict>
<key>Track ID</key><integer>%s</integer>
</dict>
"""

  filePlEndTemplate="""</array>
</dict>
"""

  plEpilogue="""</array>
</dict>
</plist>
"""
  
  # Encode stuff that needs encoding
  for trackID in tracks.keys():
    trackDict=tracks[trackID]
    dirty=0
    for trackInfoKey in kEncodeValues:
      if trackDict.has_key(trackInfoKey):
        v=trackDict[trackInfoKey]
        for chrNeedingReplacing in kEncodeChars:
          if chrNeedingReplacing in v:
            v=v.replace(chrNeedingReplacing,kEncodeCharsEncoded[chrNeedingReplacing])
            trackDict[trackInfoKey]=v
            dirty=1
    if dirty:
      tracks[trackID]=trackDict
      
  plText=filePrologue
  
  for trackID in playList["All Track IDs"]:
    trackDict=tracks[str(trackID)]
    plText+=fileTrackTemplateIDName % (trackID,trackID,trackDict["Name"])
    
    if trackDict.has_key("Artist"):
      plText+=fileTrackTemplateArtist % trackDict["Artist"]
      
    if trackDict.has_key("Album"):
      plText+=fileTrackTemplateAlbum % trackDict["Album"]
    plText+=fileTrackEnd
  
  plText+=fileMiddleTemplate # % plName
  
  for plKind in plOrder:
    if len(playList[plKind])==0:
      continue
      
    plText+=filePLStartTemplate % ("AG-"+plKind)
    
    for trackID in playList[plKind]:
      plText+=filePLTrackTemplate % trackID
    
    plText+=filePlEndTemplate
  
  plText+=plEpilogue

  return plText
  

def main():
  print "Parsing XML file..."
  e=ElementTree.parse(os.path.expanduser("~/Music/iTunes/iTunes Music Library.xml"))
  r=e.getroot()

  ur=r.getchildren()[0]
  
  topDict=dictify(ur)
  
  print
  assert topDict["Major Version"]==1
  assert topDict["Minor Version"]==1
  
  tracks=dictify(topDict["Tracks"])
    
  for k in tracks.keys()[:]:
    tracks[k]=dictify(tracks[k])
  
    # This is for an iPod, no point in adding streams
    if tracks[k]["Track Type"]=="URL":
      del tracks[k]
      continue
      
  # This shouldn't actually affect things since
  # no play date should mean a play count of 0
  # and that will get a track picked up first.
  randomOldDate=datetime.datetime(2000,1,1,0,0,0,0,utc())
  
  for track in tracks.values():
    track.setdefault("Play Count",0)
    track.setdefault("Compilation",False)
    track.setdefault("Track Number",0)
    track.setdefault("Play Date UTC",randomOldDate)
    track.setdefault("Genre","(no genre recorded)")

    # Mustn't change album, artist iTunes won't like resulting playlist
    if not track.has_key("Album"):
      track["Album for Display"]="(no album recorded)"
    else:
      track["Album for Display"]=track["Album"]

    if not track.has_key("Artist"):
      track["Artist for Display"]="(no artist recorded)"
    else:
      track["Artist for Display"]=track["Artist"]
      
    
  playList={"All Track IDs": [],
    "Recently Added": [],
    "Never Played": [],
    "Most Played": [],
    "Size": 0}
  
  for genre in kGenresToSwitchAmong:
    if genre==None:
      playList["Other Genres"]=[]
    else:
      playList[genre]=[]
  
  # Album-ize
  albumDict={}
  # Should we have a unique dummy artist per compilation and, if so,
  # how should we calculate it?
  for track in tracks.values():
    if track["Compilation"]:
      albumKey=("Dummy artist for compilation",track["Album for Display"])
    else:
      albumKey=(track["Artist for Display"],track["Album for Display"])
    albumDict.setdefault(albumKey,[])
    albumDict[albumKey].append(track)
    
  albums=albumDict.values()
  print "%i albums, %i tracks in library" % (len(albums),len(tracks))
  print
  
  try:
    for ind in range(len(albums)):
      playCounts=[]
      addedDates=[]
      playedDates=[]
      albumTracks=albums[ind]
      albumTracks.sort(key=operator.itemgetter("Track Number"))
      for track in albumTracks:
        playCounts.append(track["Play Count"])
        addedDates.append(track["Date Added"])
        playedDates.append(track["Play Date UTC"])
        
      # We use median play count for album's play count,
      # individual popular tracks will be collected later
      medianPlayCount=median(playCounts)
      totalPlayCount=sum(playCounts)
          
      addedDates.sort()
      latestAddedDate=addedDates[-1]
      # I'd like the median for last played date, but
      # this is close enough since I'm not willing to
      # go to the trouble of averaging two datetime
      # objects
      playedDates.sort()
      medianishPlayedDate=playedDates[int(len(playedDates)/2)]
        
      albums[ind]={"Tracks": albumTracks,
                   "Median Play Count": medianPlayCount,
                   "Total Play Count": totalPlayCount,
                   "Artist for Display": albums[ind][0]["Artist for Display"],
                   "Album for Display": albums[ind][0]["Album for Display"],
                   "Genre": albums[ind][0]["Genre"],
                   "Date Added": latestAddedDate,
                   "Played Date": medianishPlayedDate}
    
    # First, the albums most recently ripped
    print "Getting recently-added albums"
    added=False
    laterThan=datetime.datetime.now(localTimezone())
    laterThan-=datetime.timedelta(days=kAddAsNewDays)
    
    print "Getting albums added on or after",laterThan.astimezone(localTimezone()).strftime("%a, %b %d")
    
    albums.sort(key=operator.itemgetter("Date Added"))
    albums.reverse()
  
    for album in albums:
      if album["Date Added"]<laterThan:
        break
      
      s="%s by %s added on %s" % (album["Album for Display"],album["Artist for Display"],
        album["Date Added"].astimezone(localTimezone()).strftime("%a, %b %d"))
      print s.encode("utf-8")
  
      tryAdd(playList,album["Tracks"],"Recently Added")
      added=True
      
    if not added:
      print "None found"
    print
    
    
    # Now never-played albums
    print "Getting albums that have never been played but aren't already on the list"
    added=False
    albums.sort(key=operator.itemgetter("Total Play Count"))
    for album in albums:
      if album["Total Play Count"]>0:
        break
        
      if album["Tracks"][0]["Track ID"] not in playList["All Track IDs"]:
        s="%s by %s" % (album["Album for Display"],album["Artist for Display"])
        print s.encode("utf-8")
      
        tryAdd(playList,album["Tracks"],"Never Played")
        added=True
      
    if not added:
      print "None found"
    print


    # Now for individual tracks
    # Flatten to leave in album order
    flattenedTracks=[]
    for album in albums:
      flattenedTracks+=album["Tracks"]
    
    # Never played
    print "Getting tracks that have never been played but aren't already in the list"
    added=False
    for track in flattenedTracks:
      if track["Play Count"]==0:
        if track["Track ID"] not in playList["All Track IDs"]:
          s="%s by %s" % (track["Name"],track["Artist for Display"])
          print s.encode("utf-8")
          tryAdd(playList,[track],"Never Played")
          added=True
    
    if not added:
      print "None found"

    print

    
    # Now often-played albums
    
    print "Getting most-played albums"
    added=False
    
    albumPlayCounts=[ album["Median Play Count"] for album in albums ]
  
    stats=someStats(albumPlayCounts)
    print "Album max play count %i" % stats["max"]
    print "Album min play count %i" % stats["min"]
    print "Album mean play count %.1f" % round(stats["mean"],1)
    print "Album play count std. dev. %.1f" % round(stats["stdDev"],1)
    print "Album median play count %.1f" % round(stats["median"],1)
    
    minToGet=int(stats["max"]-stats["stdDev"]*kAlbumStdDevsToGet)
    print "Getting albums with play counts at or above %i" % minToGet
  
  
    albums.sort(key=operator.itemgetter("Median Play Count"))
    albums.reverse()
  
    for album in albums:
      if album["Median Play Count"]<minToGet:
        break
      
      # See if it's already there. We've added only albums, so
      # checking one track is sufficient.
      if album["Tracks"][0]["Track ID"] not in playList["All Track IDs"]: 
        s="%s by %s, play count %i" % (album["Album for Display"],album["Artist for Display"],album["Median Play Count"])
        print s.encode("utf-8")
      
        tryAdd(playList,album["Tracks"],"Most Played")
        added=True
        
    if not added:
      print "None found"

    print
  
  
    print "Getting most-played tracks that aren't already on the list"
    added=False
    trackPlayCounts=[ track["Play Count"] for track in flattenedTracks ]
    
    stats=someStats(trackPlayCounts)
    print "Track max play count %i" % stats["max"]
    print "Track min play count %i" % stats["min"]
    print "Track mean play count %.1f" % round(stats["mean"],1)
    print "Track play count std. dev. %.1f" % round(stats["stdDev"],1)
    print "Track median play count %.1f" % round(stats["median"],1)
    
    minToGet=int(stats["max"]-stats["stdDev"]*kTrackStdDevsToGet)
    print "Getting tracks with play counts at or above",minToGet
    
    for track in flattenedTracks:
      if track["Play Count"]>=minToGet and track["Track ID"] not in playList["All Track IDs"]:
        s="%s by %s on %s, play count %i" % (track["Name"],track["Artist for Display"],
          track["Album for Display"],track["Play Count"])
        print s.encode("utf-8")
        #tracksForPL.append(track["Track ID"])
        #sizeSoFar+=track["Size"]
        tryAdd(playList,[track],"Most Played")
        added=True

    if not added:
      print "None found"

    print
    
    # Genre-ify
    assert len(set(kGenresToSwitchAmong))==len(kGenresToSwitchAmong)
    if None in kGenresToSwitchAmong:
      dispList=list(kGenresToSwitchAmong)
      dispList.remove(None)
      dispStr=", ".join(dispList)+", (all other genres)"
    else:
      dispStr=", ".join(kGenresToSwitchAmong)
      
    print "%i genres desired: %s" %(len(kGenresToSwitchAmong),dispStr)
    print "Assembling albums by genre"
    genres={}
    for genre in kGenresToSwitchAmong:
      genres[genre]=[]
    
    for album in albums:
      g=album["Genre"]
      if g in kGenresToSwitchAmong:
        genres[g].append(album)
      elif None in kGenresToSwitchAmong:
        genres[None].append(album)
    
    for genre in genres:
      # Arrange not-recently-played mostly in front
      albumsInGenre=genres[genre]
      albumsInGenre.sort(key=operator.itemgetter("Played Date"))
      albumsInGenre=biasedSample(albumsInGenre)
      genres[genre]=albumsInGenre
  
    for genre in itertools.cycle(kGenresToSwitchAmong):
      if len(genres[genre])==0:
        continue
        
      if genre==None:
        dispStr="(all other genres)"
      else:
        dispStr=genre
      print "Adding an album from %s" % dispStr
  
      album=genres[genre].pop(0)
  
      s="%s by %s, last played on %s" % (album["Album for Display"],album["Artist for Display"],
        album["Played Date"].astimezone(localTimezone()).strftime("%a, %b %d"))
      print s.encode("utf-8")
  
      if genre==None:
        tryAdd(playList,album["Tracks"],"Other Genres")
      else:
        tryAdd(playList,album["Tracks"],genre)
  
      if len(genres[genre])==0:
        print "Genre %s exhausted" % dispStr
        if sum([ len(g) for g in genres.values()])==0:
          print "All genres exhausted"
          break
        
  except plFull:
    pass
    
    
  outputDir=os.path.expanduser(kOutputFileDir)
  outputPath=os.path.join(outputDir,kOutputFilename)
  if os.path.exists(outputPath):
    serial=1
    l=kOutputFilename.rsplit(".")
    l[0]=l[0]+"-%i"
    outputFilename=".".join(l)
    while True:
      outputPath=os.path.join(kOutputFileDir,outputFilename % serial)
      if not os.path.exists(outputPath):
        break
      else:
        serial+=1
        
  print "Creating playlist %s" % outputPath
  plText=makePLText(tracks,playList)
  
  f=open(outputPath,"w")
  
  f.write(plText.encode("utf-8"))
  f.close()
  print "Playlist written"
  return 0


if __name__=="__main__":
  main()