package com.canigraduate.uchicago.pipeline.transforms;

import com.canigraduate.uchicago.coursesearch.CourseSearch;
import com.canigraduate.uchicago.models.Course;
import com.canigraduate.uchicago.models.Term;
import com.canigraduate.uchicago.models.Watch;
import com.canigraduate.uchicago.pipeline.models.TermKey;
import org.apache.beam.sdk.transforms.FlatMapElements;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.TypeDescriptor;
import org.apache.beam.sdk.values.TypeDescriptors;

import java.util.stream.Collectors;

public class WatchToCourseTransform extends
        PTransform<PCollection<KV<String, Watch>>, PCollection<KV<TermKey, Course>>> {
    @Override
    public PCollection<KV<TermKey, Course>> expand(PCollection<KV<String, Watch>> input) {
        return input.apply(FlatMapElements
                .into(TypeDescriptors.kvs(TypeDescriptor.of(TermKey.class), TypeDescriptor.of(Course.class))).via(
                        // e is type Map<String, Watch>
                        e ->
                                CourseSearch.getCoursesFromWatch(e.getValue()).entrySet().stream()
                                        // a is type Map<String, Course>
                                        .map(a -> KV.of(TermKey.builder().setTerm(Term.create(e.getValue()
                                                        .getTerm())).setCourse(e.getValue().getCourse()).build(),
                                                a.getValue())).collect(Collectors.toList())
                ));
    }
}

@@ -676,24 +676,16 @@
 form
- extends
 %0A       
  PTr
@@ -680,16 +680,24 @@
 %0A       
+ extends
  PTransf
@@ -915,16 +915,22 @@
 Elements
+.into(
 %0A       
@@ -938,22 +938,16 @@
         
-.into(
 TypeDesc
@@ -1047,24 +1047,16 @@
         
-        
 // e is 
@@ -1099,52 +1099,12 @@
     
-        e -%3E%0A                               
+e -%3E
  Cou
@@ -1191,32 +1191,16 @@
         
- 
-               
 // a is 
@@ -1248,32 +1248,16 @@
         
-               
- 
 .map(a -
@@ -1281,16 +1281,49 @@
 uilder()
+%0A                                
 .setTerm
@@ -1347,31 +1347,29 @@
 tValue()
-%0A              
+.getTerm()))%0A
         
@@ -1396,87 +1396,31 @@
     
-          .getTerm())).setCourse(e.getValue().getCourse()).build(),%0A           
+.setCourse(a.getKey())%0A
     
@@ -1447,20 +1447,25 @@
         
-    
+.build(),
  a.getVa
@@ -1500,33 +1500,16 @@
 oList())
-%0A                
 ));%0A    

Update course in termkey.