// The mapping function is called once for each piece of the input.
// In this framework, the key is the name of the file that is being processed,
// and the value is the file's contents. The return value should be a slice of
// key/value pairs, each represented by a mapreduce.KeyValue.
func mapF(document string, value string) (res []mapreduce.KeyValue) {
  // TODO: you have to write this function
  f := func(r rune) bool {
    return !unicode.IsLetter(r)
  }
  words := strings.FieldsFunc(value, f)
  for _, word := range words {
    kv := mapreduce.KeyValue{word, " "}
    res = append(res, kv) 
  }
  return
}
 
// The reduce function is called once for each key generated by Map, with a
// list of that key's string value (merged across all inputs). The return value
// should be a single output value for that key.
func reduceF(key string, values []string) string {
  // TODO: you also have to write this function
  s := strconv.Itoa(len(values))
  return s
}
 
// Can be run in 3 ways:
// 1) Sequential (e.g., go run wc.go master sequential x1.txt .. xN.txt)
// 2) Master (e.g., go run wc.go master localhost:7777 x1.txt .. xN.txt)
// 3) Worker (e.g., go run wc.go worker localhost:7777 localhost:7778 &)
func main() {
  if len(os.Args) < 4 { 
    fmt.Printf("%s: see usage comments in file\n", os.Args[0])
  } else if os.Args[1] == "master" {
    var mr *mapreduce.Master
    if os.Args[2] == "sequential" {
      mr = mapreduce.Sequential("wcseq", os.Args[3:], 3, mapF, reduceF)
    } else {
      mr = mapreduce.Distributed("wcseq", os.Args[3:], 3, os.Args[2])
    }   
    mr.Wait()
  } else {
    mapreduce.RunWorker(os.Args[2], os.Args[3], mapF, reduceF, 100)
  }
}

// Sequential runs map and reduce tasks sequentially, waiting for each task to            // complete before scheduling the next.
func Sequential(jobName string, files []string, nreduce int,
  mapF func(string, string) []KeyValue,
  reduceF func(string, []string) string,
) (mr *Master) {
  mr = newMaster("master")
  go mr.run(jobName, files, nreduce, func(phase jobPhase) {
    switch phase {
    case mapPhase:
      for i, f := range mr.files {
        doMap(mr.jobName, i, f, mr.nReduce, mapF)
      }
    case reducePhase:
      for i := 0; i < mr.nReduce; i++ {
        doReduce(mr.jobName, i, len(mr.files), reduceF)
      }
    } 
  }, func() {
    mr.stats = []int{len(files) + nreduce}
  })  
  return
}

// doMap does the job of a map worker: it reads one of the input files
// (inFile), calls the user-defined map function (mapF) for that file's
// contents, and partitions the output into nReduce intermediate files.
func doMap(
	jobName string, // the name of the MapReduce job
	mapTaskNumber int, // which map task this is
	inFile string,
	nReduce int, // the number of reduce task that will be run ("R" in the paper)
	mapF func(file string, contents string) []KeyValue,
) {
	// TODO:
	// You will need to write this function.
	// You can find the filename for this map task's input to reduce task number
	// r using reduceName(jobName, mapTaskNumber, r). The ihash function (given
	// below doMap) should be used to decide which file a given key belongs into.
	//
	// The intermediate output of a map task is stored in the file
	// system as multiple files whose name indicates which map task produced
	// them, as well as which reduce task they are for. Coming up with a
	// scheme for how to store the key/value pairs on disk can be tricky,
	// especially when taking into account that both keys and values could
	// contain newlines, quotes, and any other character you can think of.
	//
	// One format often used for serializing data to a byte stream that the
	// other end can correctly reconstruct is JSON. You are not required to
	// use JSON, but as the output of the reduce tasks *must* be JSON,
	// familiarizing yourself with it here may prove useful. You can write
	// out a data structure as a JSON string to a file using the commented
	// code below. The corresponding decoding functions can be found in
	// common_reduce.go.
	//
	//   enc := json.NewEncoder(file)
	//   for _, kv := ... {
	//     err := enc.Encode(&kv)
	//
	// Remember to close the file after you have written all the values!``
	file, err := os.Open(inFile)
	if err != nil {
		log.Fatal(err)
	}
	defer file.Close()
	scanner := bufio.NewScanner(file)
	contents := ""
	for scanner.Scan() {
		s := scanner.Text()
		s += "\n"
		contents = contents + s
	}
	kvs := mapF(inFile, contents)
	reduceFileMap := make(map[string]*os.File)
	jsonFileMap := make(map[string]*json.Encoder)
	for i := 0; i < nReduce; i++ {
		reduceFileName := reduceName(jobName, mapTaskNumber, i)
		file, err := os.Create(reduceFileName)
		if err != nil {
			log.Fatal(err)
		}
		enc := json.NewEncoder(file)
		reduceFileMap[reduceFileName] = file
		jsonFileMap[reduceFileName] = enc
		defer reduceFileMap[reduceFileName].Close()
	}
	for _, kv := range kvs {
		hashValue := int(ihash(kv.Key))
		fileNum := hashValue % nReduce
		reduceFileName := reduceName(jobName, mapTaskNumber, fileNum)
		enc, ok := jsonFileMap[reduceFileName]
		if !ok {
			log.Fatal(err)
		}
		err := enc.Encode(&kv)
		if err != nil {
			log.Fatal(err)
		}
	}
}

// doReduce does the job of a reduce worker: it reads the intermediate
// key/value pairs (produced by the map phase) for this task, sorts the
// intermediate key/value pairs by key, calls the user-defined reduce function
// (reduceF) for each key, and writes the output to disk.
func doReduce(
	jobName string, // the name of the whole MapReduce job
	reduceTaskNumber int, // which reduce task this is
	nMap int, // the number of map tasks that were run ("M" in the paper)
	reduceF func(key string, values []string) string,
) {
	// TODO:
	// You will need to write this function.
	// You can find the intermediate file for this reduce task from map task number
	// m using reduceName(jobName, m, reduceTaskNumber).
	// Remember that you've encoded the values in the intermediate files, so you
	// will need to decode them. If you chose to use JSON, you can read out
	// multiple decoded values by creating a decoder, and then repeatedly calling
	// .Decode() on it until Decode() returns an error.
	//
	// You should write the reduced output in as JSON encoded KeyValue
	// objects to a file named mergeName(jobName, reduceTaskNumber). We require
	// you to use JSON here because that is what the merger than combines the
	// output from all the reduce tasks expects. There is nothing "special" about
	// JSON -- it is just the marshalling format we chose to use. It will look
	// something like this:
	//
	// enc := json.NewEncoder(mergeFile)
	// for key in ... {
	// 	enc.Encode(KeyValue{key, reduceF(...)})
	// }
	// file.Close()
	kvs := make(map[string][]string)
	for i := 0; i < nMap; i++ {
		reduceFileName := reduceName(jobName, i, reduceTaskNumber)
		file, err := os.Open(reduceFileName)
		if err != nil {
			log.Fatal(err)
		}
		defer file.Close()
		enc := json.NewDecoder(file)
		for {
			var kv KeyValue
			if err := enc.Decode(&kv); err == io.EOF {
				break
			} else if err != nil {
				log.Fatal(err)
			} else {
				log.Println(kv.Key + kv.Value)
				kvs[kv.Key] = append(kvs[kv.Key], kv.Value)
			}
		}
	}

	var keys []string
	for k, _ := range kvs {
		keys = append(keys, k)
	}
	sort.Sort(sort.StringSlice(keys))
	mergeFileName := mergeName(jobName, reduceTaskNumber)
	mergeFile, err := os.Create(mergeFileName)
	if err != nil {
		log.Fatal(err)
	}
	defer mergeFile.Close()
	enc := json.NewEncoder(mergeFile)
	for _, key := range keys {
		enc.Encode(KeyValue{key, reduceF(key, kvs[key])})
	}
}

// schedule starts and waits for all tasks in the given phase (Map or Reduce).
func (mr *Master) schedule(phase jobPhase) {
	var ntasks int
	var nios int // number of inputs (for reduce) or outputs (for map)
	switch phase {
	case mapPhase:
		ntasks = len(mr.files)
		nios = mr.nReduce
	case reducePhase:
		ntasks = mr.nReduce
		nios = len(mr.files)
	}

	fmt.Printf("Schedule: %v %v tasks (%d I/Os)\n", ntasks, phase, nios)

	// All ntasks tasks have to be scheduled on workers, and only once all of
	// them have been completed successfully should the function return.
	// Remember that workers may fail, and that any given worker may finish
	// multiple tasks.
	//
	// TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO
	//
	switch phase {
	case mapPhase:
		mr.scheduleMap()
	case reducePhase:
		mr.scheduleReduce()
	}
	fmt.Printf("Schedule: %v phase done\n", phase)
}