accession   geneName  specie  sequence 
P01308	    INS	      HUMAN   MALWMRLLPLLALLALWGPDPAAAFVNQHLCGSHLVEALYLVCGERGFFYTPKTRREAEDLQVGQVELGGGPGAGSLQPLALEGSLQKRGIVEQCCTSICSLYQLENYCN

//Read data first (you need to uncompress the file in this case)
val data = sc.textFile("swissprot-aug-2017.tsv")
//data: org.apache.spark.rdd.RDD[String] = ...

//Or if working with gz file
val data = sc.textFile("swissprot-aug-2017.tsv.gz")

//Shows the first line
data.first 

//Defines a class to represent the data in TSV 
case class Sequence(accession: String, geneName: String, specie: String, sequence: String)

//Reads a line from the TSV file and returns an object from the Sequence class
def readLine(line: String) = {
  val parts = line.split("\t"); 
  Sequence(parts(0), parts(1), parts(2), parts(3))
}

//In this case map tells how each line should be mapped / transformed (nothing is actually executed)
val sequences = data.map(readLine) //transformation 1

//Shows first accession or specie (this is an action, but does not require to read the whole file)
sequences.first // action
sequences.first.getClass //action: Notice that this is a class and not an RDD

sequences.first.accession //action
sequences.first.specie //action

//If you have initialized the spark-shell with 2G you may be able to store the result in cache (big difference with Hadoop)
sequences.cache //Asks Spark to persist result in memory if possible

//See how many different species we have
sequences.map(s => s.specie).distinct().count //2 transformations and 1 action

//See the species with the most number of sequences
//3 transformations and 1 action
sequences.map(s => (s.specie, 1)).reduceByKey(_ + _).sortBy(-_._2).take(10).foreach(println)

// This is an action, but all rows must be accessed, because file is not indexed yet.
val humanSequences = sequences.filter(s=> s.specie.equals("HUMAN")) //1 transformation
//Count the number of human sequences
humanSequences.count //1 action

//Sequences containing the 'DANIEL' sequence 
val sequencesContainingPattern = sequences.filter(s=> s.sequence.contains("DANIEL")).count

//Use any function (alignment, GC content, ..... ) 
val result = sequences.map(ANY_FUNCTION_HERE)

object Enzyme extends java.io.Serializable {
  //Suppose the enzyme recognizes the portion of the following sequence
  val cleavageSite = """[ASD]{3,5}[LI][^P]{2,5}""".r
  def matchSite(seq: String): Boolean = cleavageSite.findFirstIn(seq).isDefined
}

//Searches for all sequences that the enzyme will recognize
sequences.filter(s => Enzyme.matchSite(s.sequence)).count

//Another example with sorting
object PrositePatterns extends java.io.Serializable {
  //Glycosylation prosite pattern
  val nglycoPattern = """N[^P][ST][^P]""".r
  def countNGlycoSites (seq: String) : Integer = nglycoPattern.findAllIn(seq).length
}

//Returns sequences that have more than 100 N-glycosylation sites predicted (this prediction returns many false positives)
sequences.map(s => (s.accession, s.geneName, PrositePatterns.countNGlycoSites(s.sequence))).filter(r => r._3 > 100).sortBy(-_._3).take(10)

case class Peptide (accession: String, startPosition: Integer, length: Integer)

object Trypsin extends java.io.Serializable{
  //The trypsin cuts just after Arginine (R) or Lysine (K), but not with a Proline (P) after
  val pattern = """(?!P)(?<=[RK])""".r
  def cut (accession: String, sequence: String) : List[Peptide] = {
    val startPositions : List[Int] = List(0) ++ pattern.findAllMatchIn(sequence).map(_.start).toList ++ List(sequence.length)
    startPositions.sliding(2).map(l => Peptide(accession, l(0) + 1, l(1)-l(0))).toList
  }
}

//Check the top 10 biggest peptides for human sequences
humanSequences.flatMap(s => Trypsin.cut(s.accession, s.sequence)).sortBy(s => -s.length).take(10)

//Draw peptide frequency by peptide length
humanSequences.flatMap(s => Trypsin.cut(s.accession, s.sequence)).map(pep => pep.length).countByValue.toList.sortBy(_._1).foreach(p => println(p._1 + "\t" + p._2 ))

//Need to add --packages at the end (corresponding to Java/Scala artifact. Downloaded by sbt automatically and deployed to worker nodes.
$SPARK_HOME/bin/spark-shell --master spark://$master_hostname:7077 --packages "com.fulcrumgenomics:fgbio_2.11:0.2.0"

import com.fulcrumgenomics.alignment.NeedlemanWunschAligner

val seqToAlign = "MALWMRLLPLLALLALWGPDPAAAFVNQHLCGSHLVEALYLVCGERGFFYTPKTRREAEDLQVGQVELGGGPGAGSLQPLALEGSLQKRGIVEQCCTSICSLYQLENYCD"

humanSequences.map(s=> (s.accession, NeedlemanWunschAligner(1, -1, -3, -1).align(seqToAlign, s.sequence).score)).sortBy(-_._2).take(10).foreach(println)

//The use of data frames may require to import SQL context
val sqlContext = new org.apache.spark.sql.SQLContext(sc)
import sqlContext.implicits._

//Transforms the RDD to a DataFrame
val df = sequences.toDF
df.show
// Print the schema in a tree format (you can imagine nested classes inside the Sequence class)
df.printSchema()

df.createOrReplaceTempView("sequences")

val sqlQuery = """SELECT specie, count(*) as cnt 
FROM sequences 
GROUP BY specie 
ORDER BY cnt DESC 
LIMIT 10"""

sqlContext.sql(sqlQuery).show

//OR directly 
df.groupBy("specie").agg(count("*") as "cnt").orderBy(desc("cnt")).limit(10).show

//Or a more complex example, that also computes the average of length of the sequence
//This query sorts by the average length, you can see the sequence with less amino acids:
//http://www.uniprot.org/uniprot/P84761

val sqlQuery2 = """SELECT specie, count(*) as cnt, avg(length(sequence)) as avgSeqLength 
FROM sequences 
GROUP BY specie 
ORDER BY avgSeqLength ASC LIMIT 10"""

//TSV file: swissprot-aug-2017.tsv (200MB)
//Compressed file: swissprot-aug-2017.tsv.gz (66MB)
//Let's see how much space is the equivalent in Parquet
val data = sc.textFile("swissprot-aug-2017.tsv")
case class Sequence(accession: String, geneName: String, specie: String, sequence: String)
def readLine(line: String) = {
  val parts = line.split("\t"); 
  Sequence(parts(0), parts(1), parts(2), parts(3))
}
val sequences = data.map(readLine) 

//Save into format Parquet format
//swissprot-aug-2017.parquet folder (109 MB)
sequences.toDF.write.format("parquet").save("swissprot-aug-2017.parquet")

val sqlContext = new org.apache.spark.sql.SQLContext(sc)
import sqlContext.implicits._
val sequences = sqlContext.read.parquet("swissprot-aug-2017.parquet")
sequences.printSchema()
//Schema is already defined in the parquet format. No need to define class 
sequences.groupBy($"specie").agg(count("*") as "cnt").orderBy($"cnt" desc).limit(10).show

val data = sc.textFile("swissprot-aug-2017.tsv")
val sequences = data.map(readLine)
val humanSequences = sequences.filter(s=> s.specie.equals("HUMAN"))
val aaFrequency = humanSequences.flatMap(s=> s.sequence.toList).map(aa => (aa, 1L)).reduceByKey(_ + _).sortBy(-_._2)

//The collect or take action will trigger the rendering of a table that can be toggled to a chart
aaFrequency.collect

import sys
from Bio import SeqIO
fasta_file = sys.argv[1]
for seq_record in SeqIO.parse(fasta_file, "fasta"):
    header = str(seq_record.id).split('|')
    accessionSpecie = str(header[2]).split("_")
    sys.stdout.write(str(header[1])+'\t'+str(accessionSpecie[0])+'\t'+str(accessionSpecie[1])+'\t'+str(seq_record.seq)+ '\n')